Source-linked AI summary
Mask-guided Spectral-wise Transformer for Efficient Hyperspectral Image Reconstruction
Yuanhao Cai, Jing Lin, Xiaowan Hu, Haoqian Wang, Xin Yuan, Yulun Zhang, Radu Timofte, Luc Van Gool
TL;DR
HSI reconstruction must recover a 3D spatial-spectral signal from a 2D CASSI measurement, while existing CNNs inadequately model spectral correlations, long-range dependencies, and mask guidance. MST uses spectral-wise self-attention and a mask-guided mechanism to target high-fidelity regions. It outperforms SOTA methods on simulation and real HSI data with lower computational and memory costs.
Problem
Existing CNN-based HSI reconstruction methods have limitations in modeling spectral-wise similarity and long-range dependencies, and do not fully exploit CASSI mask guidance.
Method
MST treats spectral feature channels as tokens for S-MSA and uses a mask-guided mechanism to focus attention on regions with high-fidelity spectral representations.
Results
MST dramatically outperforms SOTA methods on all simulated scenes while requiring much cheaper Params and FLOPS, and produces more visually pleasant real-world reconstructions.
Takeaways & Limitations
MST provides an efficient Transformer-based approach for HSI reconstruction that models inter-spectral dependencies and uses physical-mask information.
Abstract
from arXiv · showhide
Hyperspectral image (HSI) reconstruction aims to recover the 3D spatial-spectral signal from a 2D measurement in the coded aperture snapshot spectral imaging (CASSI) system. The HSI representations are highly similar and correlated across the spectral dimension. Modeling the inter-spectra interactions is beneficial for HSI reconstruction. However, existing CNN-based methods show limitations in capturing spectral-wise similarity and long-range dependencies. Besides, the HSI information is modulated by a coded aperture (physical mask) in CASSI. Nonetheless, current algorithms have not fully explored the guidance effect of the mask for HSI restoration. In this paper, we propose a novel framework, Mask-guided Spectral-wise Transformer (MST), for HSI reconstruction. Specifically, we present a Spectral-wise Multi-head Self-Attention (S-MSA) that treats each spectral feature as a token and calculates self-attention along the spectral dimension. In addition, we customize a Mask-guided Mechanism (MM) that directs S-MSA to pay attention to spatial regions with high-fidelity spectral representations. Extensive experiments show that our MST significantly outperforms state-of-the-art (SOTA) methods on simulation and real HSI datasets while requiring dramatically cheaper computational and memory costs. Code and pre-trained models are available at https://github.com/caiyuanhao1998/MST/
1. Introduction
HSI reconstruction recovers a 3D spatial-spectral signal from a 2D CASSI measurement, but CNN methods inadequately model spectral similarity, long-range dependencies, and mask guidance. MST addresses these issues with spectral-wise attention and a mask-guided mechanism, achieving strong reconstruction results with lower costs.
- HSIs provide wavelength-specific channels that store richer scene information than RGB images.
- CASSI compresses snapshot information along the spectral dimension into a single 2D measurement for HSI capture.
- CNN-based reconstruction methods are limited in modeling inter-spectra similarity and long-range dependencies, while prior mask usage mainly combines the mask with shifted measurements.
- MST uses S-MSA to model long-range inter-spectra dependencies by treating each spectral channel feature as a token.
- MST adds a mask-guided mechanism that directs attention toward regions with high-fidelity spectral representations.
- MST dramatically outperforms SOTA methods on all simulated scenes while using much cheaper Params and FLOPS, and produces more visually pleasant real-world reconstructions.
2. Related Work
Related work spans model-based priors, CNN reconstruction, and Transformers, but prior approaches do not fully address spectrally correlated, mask-modulated HSI data.
- Model-based HSI reconstruction uses hand-crafted priors such as sparsity, total variation, and non-local similarity.
- Deep CNNs learn end-to-end mappings for HSI reconstruction and have achieved promising performance.
- Vision Transformers capture spatial long-range dependencies across high-level and low-level vision tasks, including image restoration.
- Directly applying existing Transformers to HSIs is unsuitable because HSI representations are spectrally self-similar and modulated by a physical mask.
3. CASSI System
CASSI modulates and spectrally shears a 3D HSI cube before summing its channels into a noisy 2D detector measurement.
- A coded aperture modulates every spectral channel of the HSI through element-wise multiplication with a physical mask.
- The disperser shears each modulated spectral channel along the detector’s y-axis according to its wavelength relative to a reference channel.
- The detector forms the compressed measurement by summing all sheared spectral channels and adding imaging noise.
4. Method
MST combines a U-shaped reconstruction architecture with mask-guided spectral attention, modeling inter-spectral dependencies while retaining efficient hierarchical processing.
- Overall Architecture: MST reverses CASSI dispersion by shifting the measurement back to initialize an HSI signal before feature extraction.
- Overall Architecture: Its encoder, bottleneck, and decoder use hierarchical features, downsampling, upsampling, and skip connections to reconstruct residual HSIs.
- Overall Architecture: MSAB combines layer normalization, MS-MSA, and an FFN as MST’s basic processing unit.
- Spectral-wise Multi-head Self-Attention: S-MSA treats each spectral feature map as a token and computes self-attention along the spectral dimension rather than spatially.
- Spectral-wise Multi-head Self-Attention: The attention projections use learnable query, key, and value parameters and split features into multiple spectral heads.
- Spectral-wise Multi-head Self-Attention: A learnable parameter σ_j adapts each head’s attention by reweighting K_j^T Q_j to accommodate wavelength-dependent spectral density.
S - MSA}(\
S-MSA models inter-spectral dependencies by treating spectral feature maps as tokens, while MM uses position-sensitive mask fidelity to guide spatial attention. The design retains global receptive fields with spatially efficient computation.
- S-MSA: The attention output concatenates N heads, applies a learnable linear projection, and adds a position embedding before reshaping to feature maps.
- S-MSA: S-MSA and W-MSA have computational complexity linear in spatial size HW, whereas G-MSA is quadratic in HW.
- Mask-guided Mechanism: The mask is shifted according to wavelength-dependent dispersion so its attention map aligns with modulated HSI features.
- Mask-guided Mechanism: MM processes shifted mask features through parallel identity and spatial-correlation paths to produce guidance for S-MSA.
- Mask-guided Mechanism: MM re-weights each attention head's value features with mask guidance, enabling extraction of non-corrupted representations and adaptive spatial interactions.
5. Experiments
Experiments evaluate MST on simulation and real HSI reconstruction, comparing quality, efficiency, visual fidelity, and ablations of S-MSA and MM. MST outperforms prior methods while using lower computational and memory costs, and both proposed components improve reconstruction.
- Experimental Settings: MST is evaluated on simulation and real HSI datasets using PSNR and SSIM, with CAVE for training and 10 KAIST scenes for testing.The real dataset was collected using a CASSI system.
- Quantitative Results: MST significantly surpasses prior methods across all 10 simulation scenes while requiring cheaper memory and computational costs.MST-L exceeds DGSMP, TSA-Net, and λ-net by 2.55, 3.72, and 6.65 dB, respectively.
- Quantitative Results: MST occupies the upper-left region of PSNR–Params–FLOPS comparisons, indicating high reconstruction performance with lower computational and parameter costs.The comparison uses PSNR for performance, FLOPS for computational cost, and circle radius for Params.
- Qualitative Results: MST-L produces sharper simulation reconstructions, preserves homogeneous-region smoothness, and more closely matches ground-truth spectral curves.The comparison uses Scene 5 with 4 of 28 spectral channels.
- Qualitative Results: MST-L reconstructs more high-frequency structural detail and suppresses real measurement noise more effectively than previous algorithms.The real-data comparison uses Scene 3 with 4 of 28 spectral channels.
- Ablation Study: S-MSA and MM improve the baseline by 0.89 and 1.08 dB, while S-MSA yields the largest gain among compared attention schemes with the lowest added FLOPS.MM also directs attention toward detailed, high-fidelity structural features in visualizations.
6. Conclusion
MST is presented as an efficient Transformer-based framework for accurate HSI reconstruction, combining spectral dependency modeling with mask-guided attention. Experiments report better reconstruction quality, visual appearance, and efficiency than SOTA methods.
- MST combines S-MSA for inter-spectra similarity and dependencies with MM guidance toward high-fidelity spatial regions.S-MSA models spectral interactions, while MM uses the mask to guide attention.
- Quantitative experiments show MST surpasses SOTA algorithms by a large margin while requiring significantly cheaper Params and FLOPS.
- Qualitative comparisons show MST produces more visually pleasant reconstructed HSIs.