[Paper Review] S^2-Transformer for Mask-Aware Hyperspectral Image Reconstruction

Jiamian Wang, Kunpeng Li+3 more|arXiv (Cornell University)|Sep 24, 2022

Photoacoustic and Ultrasonic Imaging4 citations

TL;DR

This paper proposes an S²-Transformer network with mask-aware learning for high-fidelity hyperspectral image reconstruction from CASSI measurements. By jointly modeling spatial and spectral attention and adaptively prioritizing reconstruction loss in masked regions, the method achieves state-of-the-art performance, improving PSNR by 0.17 dB and SSIM by 0.0015 over prior methods.

ABSTRACT

Snapshot compressive imaging (SCI) surges as a novel way of capturing hyperspectral images. It operates an optical encoder to compress the 3D data into a 2D measurement and adopts a software decoder for the signal reconstruction. Recently, a representative SCI set-up of coded aperture snapshot compressive imager (CASSI) with Transformer reconstruction backend remarks high-fidelity sensing performance. However, dominant spatial and spectral attention designs show limitations in hyperspectral modeling. The spatial attention values describe the inter-pixel correlation but overlook the across-spectra variation within each pixel. The spectral attention size is unscalable to the token spatial size and thus bottlenecks information allocation. Besides, CASSI entangles the spatial and spectral information into a 2D measurement, placing a barrier for information disentanglement and modeling. In addition, CASSI blocks the light with a physical binary mask, yielding the masked data loss. To tackle above challenges, we propose a spatial-spectral (S2-) Transformer implemented by a paralleled attention design and a mask-aware learning strategy. Firstly, we systematically explore pros and cons of different spatial (-spectral) attention designs, based on which we find performing both attentions in parallel well disentangles and models the blended information. Secondly, the masked pixels induce higher prediction difficulty and should be treated differently from unmasked ones. We adaptively prioritize the loss penalty attributing to the mask structure by referring to the mask-encoded prediction as an uncertainty estimator. We theoretically discuss the distinct convergence tendencies between masked/unmasked regions of the proposed learning strategy. Extensive experiments demonstrate that on average, the results of the proposed method are superior over the state-of-the-art method.

Motivation & Objective

• To address two major data losses in CASSI-based hyperspectral imaging: entangled spectral channel blending and masked pixel loss due to physical coded apertures.
• To improve reconstruction fidelity by leveraging inherent spatial and spectral correlations in hyperspectral data through a novel Transformer architecture.
• To model pixel-wise reconstruction difficulty caused by optical masking, distinguishing it from texture-induced uncertainty.
• To enhance interpretability of deep learning models for HSI reconstruction by grounding design in physical imaging principles.

Proposed method

• Proposes an S²-Transformer with parallel spatial and spectral self-attention mechanisms to disentangle entangled spectral and spatial information in 2D CASSI measurements.
• Introduces a mask-encoding (ME) loss term that penalizes predictions in early network layers based on the physical mask pattern.
• Develops a mask-aware (MA) loss that dynamically infers pixel-wise reconstruction difficulty from predicted outputs and prioritizes masked regions with higher loss weight.
• Employs a dual-branch architecture where spatial and spectral attention are processed in parallel and interact through feature concatenation and residual connections.
• Uses a pre-training stage with ME loss to stabilize early training and improve convergence in masked regions.
• Theoretical analysis shows distinct convergence behaviors between masked and unmasked regions under the proposed loss strategy.

Experimental results

Research questions

• RQ1 How can spatial and spectral attention mechanisms be jointly designed to better exploit hyperspectral data characteristics in 2D compressed measurements?
• RQ2 To what extent does modeling optical-induced data uncertainty—specifically from physical masking—improve reconstruction fidelity compared to standard L1/L2 losses?
• RQ3 Can a dynamic, mask-aware loss strategy that adapts to prediction confidence improve convergence and performance in masked regions?
• RQ4 How does the interaction between spatial and spectral attention affect the representational capacity and interpretability of the model?

Key findings

• The proposed S²-Transformer with mask-aware learning achieves a PSNR of 36.48 dB and SSIM of 0.9584 on benchmark datasets, outperforming previous state-of-the-art methods by 0.17 dB in PSNR and 0.0015 in SSIM.
• Ablation studies confirm that combining both ME and MA losses yields the best performance, while using either loss alone provides minimal or no improvement.
• The mask-aware (MA) loss significantly reduces absolute reconstruction error in masked regions, as shown by convergence curves and visual comparisons.
• Visual results demonstrate that mask-aware learning enhances texture recovery and fine detail reconstruction, especially in masked areas, both in simulation and real-world data.
• Theoretical and empirical analysis confirms that the proposed loss strategy leads to faster and more stable convergence in masked regions compared to unmasked ones.
• The parallel spatial-spectral attention mechanism effectively captures long-range spectral dependencies and spatial correlations, as evidenced by the correlation matrix analysis and attention visualization.