MS²A-Net: Memory-Augmented Spatial-Spectral Attention-Based Deep Unfolding Network for Hyperspectral Image Reconstruction

Hyperspectral snapshot compressive imaging enables high-resolution spectral acquisition in a single exposure. However, accurate reconstruction from coded measurements remains challenging because of model mismatch, measurement noise, and the need to jointly exploit spatial structure and spectral correlations. Moreover, although recent deep unfolding networks have begun to exploit interstage interactions, many CASSI unfolding architectures still pass only image estimates between stages or employ comparatively heavy feature aggregation, leaving a need for selective and cost-controlled cross-stage refinement. To address these issues, we propose MS²A-Net, a framework that integrates cross-stage memory propagation with linear-complexity spatial-spectral attention to achieve high-fidelity reconstruction. Specifically, we unfold the generalized alternating projection (GAP) algorithm and adopt a learnable compensation branch to refine the GAP-based reconstruction update using the measurement and sensing mask. In the denoising stage, to reduce the computational burden of attention, we design a memory-augmented spatial-spectral attention module (MS²AM). This module couples agent attention with windowed spectral self-attention to model complementary spatial contexts and spectral signatures with linear complexity. Furthermore, a memory-augmented gating unit (MAGU) propagates compact hidden states across unfolding stages, adaptively fuses newly extracted features with stored information, and mitigates stage-wise information loss. Experiments on simulated data demonstrate strong quantitative performance in terms of PSNR and SSIM, while qualitative evaluations on real CASSI measurements show visually coherent structures and reduced artifacts. The code is available at https://github.com/chenx2000/MS2A-Net.