Source-linked AI summary
Coarse-to-Fine Sparse Transformer for Hyperspectral Image Reconstruction
Yuanhao Cai, Jing Lin, Xiaowan Hu, Haoqian Wang, Xin Yuan, Yulun Zhang, Radu Timofte, Luc Van Gool
TL;DR
CASSI reconstruction must recover a 3D HSI cube from a single 2D measurement while existing CNNs and Transformers inadequately handle long-range dependencies, non-local similarity, and spatial sparsity. CST screens informative patches with SASM and clusters related pixels with SAH-MSA. The method significantly outperforms state-of-the-art methods at lower computational cost and reports more detailed real-restoration results.
Problem
Existing CNNs have limited non-local and long-range modeling, while Transformers process uninformative or unrelated tokens despite HSI spatial sparsity.
Method
CST uses SASM for coarse informative-patch selection and SAH-MSA for fine clustering of closely related tokens before self-attention.
Results
CST significantly outperforms state-of-the-art methods with cheaper computational costs and more visually pleasing, fine-grained real HSI reconstruction.
Takeaways & Limitations
Embedding HSI sparsity into the learning pipeline provides CST with an efficient coarse-to-fine reconstruction strategy within the reported experiments.
Takeaways & Limitations
SASM supervision assumes that dark, uninformative regions correspond approximately to regions that are hard to reconstruct.
Abstract
from arXiv · showhide
Many algorithms have been developed to solve the inverse problem of coded aperture snapshot spectral imaging (CASSI), i.e., recovering the 3D hyperspectral images (HSIs) from a 2D compressive measurement. In recent years, learning-based methods have demonstrated promising performance and dominated the mainstream research direction. However, existing CNN-based methods show limitations in capturing long-range dependencies and non-local self-similarity. Previous Transformer-based methods densely sample tokens, some of which are uninformative, and calculate the multi-head self-attention (MSA) between some tokens that are unrelated in content. This does not fit the spatially sparse nature of HSI signals and limits the model scalability. In this paper, we propose a novel Transformer-based method, coarse-to-fine sparse Transformer (CST), firstly embedding HSI sparsity into deep learning for HSI reconstruction. In particular, CST uses our proposed spectra-aware screening mechanism (SASM) for coarse patch selecting. Then the selected patches are fed into our customized spectra-aggregation hashing multi-head self-attention (SAH-MSA) for fine pixel clustering and self-similarity capturing. Comprehensive experiments show that our CST significantly outperforms state-of-the-art methods while requiring cheaper computational costs. The code and models will be released at https://github.com/caiyuanhao1998/MST
1 Introduction
CASSI reconstructs a 3D HSI cube from a single 2D compressive measurement, but existing CNN and Transformer methods inadequately exploit long-range dependencies, content similarity, and spatial sparsity. CST addresses these issues by screening informative patches and clustering related tokens for efficient self-attention.
- Motivation: CASSI compresses a 3D HSI cube into one 2D measurement for subsequent reconstruction.A coded aperture modulates wavelengths, a disperser shifts them spatially, and reconstruction restores the cube.
- Motivation: CNN-based reconstruction methods struggle to capture non-local self-similarity and long-range dependencies.
- Motivation: Existing Transformers process uninformative spatial tokens and unrelated content, reducing cost-effectiveness and potentially producing over-smooth results.Global Transformer complexity is quadratic in spatial dimensions.
- Proposed Method: CST introduces spectra-aware screening mechanism (SASM) for informative patch selection and spectra-aggregation hashing multi-head self-attention (SAH-MSA) for related-token clustering.Only detected patches enter attention, while closely related tokens are grouped into buckets.
- Results: CST significantly surpasses state-of-the-art methods with lower computational complexity and produces more visually pleasant real-restoration results.
2 Related Work
Prior HSI reconstruction methods include hand-crafted-prior approaches, CNN architectures, and Transformer-based models. These families differ in speed, generalization, and ability to model long-range interactions.
- Traditional Methods: Traditional methods use hand-crafted priors such as gradients, total variation, low rank, and non-local self-similarity.
- Traditional Methods: Model-based methods suffer from low reconstruction speed and poor generalization ability.
- CNN Methods: CNN-based algorithms include end-to-end, deep unfolding, and plug-and-play approaches for learning HSI restoration mappings.
- Transformer Methods: Transformers have been adopted in vision because they model long-range interactions between spatial regions.
3 Mathematical Model of CASSI
CASSI modulates and spatially shifts spectral channels before integrating them into a noisy 2D detector measurement. Reconstruction then seeks to recover the original 3D HSI cube from that measurement.
- Modulation: A coded aperture M* modulates each spectral channel of the 3D HSI cube F through an element-wise product.
- Dispersion: After modulation, a disperser shears each wavelength channel by an offset relative to a reference wavelength.
- Measurement Formation: The shifted spectral channels are integrated over wavelengths and combined with random imaging noise G to form the 2D measurement Y.
- Inverse Problem: Given Y captured by CASSI, reconstruction restores the original 3D HSI data cube F.
4 Method
CST reconstructs HSIs through a coarse-to-fine sparse Transformer that first selects informative patches and then restricts attention to content-related pixels. Its architecture combines SASM with SAH-MSA to reduce inefficient computation while modeling spectral and spatial dependencies.
- Network Architecture: CST comprises SASM for coarse patch selection and SAH-MSA for fine pixel clustering within spectra-aware hashing attention blocks.The three-stage symmetric encoder-decoder uses SAHAB units, skip aggregation, and residual HSI prediction.
- Spectra-Aware Screening Mechanism: SASM partitions feature maps into non-overlapping M × M patches and screens patches using a predicted sparsity mask before self-attention.Average pooling and binarization produce a patch mask, while the sparsity ratio σ controls the selected top-k patches.
- Spectra-Aware Screening Mechanism: The sparsity estimator predicts regions containing HSI representations and is trained with a reference mask derived from spectral reconstruction differences.The reference assumes dark background regions are relatively uninformative and approximately correspond to hard-to-reconstruct regions.
- Sparsity Loss: The sparsity loss is combined with the HSI reconstruction loss as a weighted objective controlled by λ.The loss balances reconstruction error with the mean squared error between predicted and reference sparsity masks.
- Spectra-Aggregation Hashing Multi-head Self-Attention: SAH-MSA hashes spectral pixel vectors, sorts them by hash codes, groups them into buckets, and computes attention only among closely correlated tokens.Multiple hashing rounds reduce the chance that similar elements fall into different buckets.
5 Experiment
Experiments use simulation and real CASSI settings with standard reconstruction metrics and visual comparisons. CST achieves stronger reconstruction than competing methods while using substantially fewer computational resources.
- Experiment Settings: Simulation experiments use CAVE for training and 10 KAIST scenes for testing, with 28 interpolated wavelengths spanning 450 nm to 650 nm.PSNR and SSIM are used to evaluate HSI reconstruction.
- Qualitative Results: On real HSI reconstruction with injected 11-bit shot noise, CST-L shows advantages in fine-grained content restoration and noise removal.The authors report these results as evidence of robustness, reliability, and generalization ability.
- Quantitative Results: CST-S, CST-M, and CST-L improve over MST by 0.45, 0.37, and 0.67 dB while using 1.29G, 1.16G, and 0.34G fewer FLOPS, respectively.These comparisons quantify the trade-off between reconstruction quality and computational cost against the recent Transformer-based MST.
- Quantitative Results: CST-L outperforms DGSMP, TSA-Net, and λ-Net by 3.22, 4.39, and 7.32 dB while using lower parameter and FLOP budgets.The reported reductions include 4.3% FLOPS versus DGSMP, 25.3% versus TSA-Net, and 23.6% versus λ-Net.
5.3 Qualitative Results
CST produces sharper, more visually pleasing simulation reconstructions while preserving smooth homogeneous regions and spectral consistency. On real HSI data, CST-L also restores fine-grained content and removes noise effectively.
- Simulation HSI Restoration: CST restores sharper edge details and smooth homogeneous regions without introducing artifacts in simulated Scene 2 reconstructions.Other methods either oversmooth structural and high-frequency details or produce blotchy textures and chromatic artifacts.
- Simulation HSI Restoration: CST achieves the highest correlation with the ground-truth spectral density curve in the evaluated region.The comparison supports spectral-dimension consistency in the reconstructed HSI.
- Real HSI Restoration: CST-L shows advantages in fine-grained content restoration and real noise removal on real HSI reconstruction.The evaluation retrains CST-L with shot noise injected into measurements from the KAIST and CAVE datasets.
5.4 Ablation Study
Ablations show that SAH-MSA provides most of the reconstruction gain, while SASM substantially reduces computation with minimal performance loss. SASM also outperforms random and uniform sparsity schemes.
- Break-down Ablation: Adding SAH-MSA improves the baseline by 2.96 dB in PSNR and 0.042 in SSIM.The baseline reaches 32.57 dB PSNR and 0.906 SSIM before adding SAH-MSA.
- Break-down Ablation: Adding SASM reduces computational cost by 31.3% while degrading PSNR by 0.6% and SSIM by 0.1%.The reported cost reduction is 7.69 / 24.60, indicating a large efficiency gain for a small reconstruction trade-off.
- Sparsity Mechanisms: SASM outperforms random and uniform sparsity schemes by over 0.9 dB and responds more completely to informative HSI regions.Random and uniform schemes select preset positions without awareness of HSI signals.
Sparsity Scheme Comparison.
SASM selects spatial patches according to HSI informativeness, unlike fixed random or uniform schemes. SAH-MSA further focuses attention on content-related tokens, improving computational cost-effectiveness.
- Sparsity Scheme Comparison: SASM generates more complete and accurate responses to informative HSI regions than random and uniform sparsity masks.Yellow marks selected patches, while green marks unselected patches in the visualization.
- Sparsity Scheme Comparison: Random and uniform sparsity schemes rigidly select preset positions without awareness of HSI signals.This contrasts with SASM’s signal-aware localization of regions containing dense HSI representations.
- Self-Attention Mechanism Comparison: SAH-MSA applies attention between closely related tokens within each bucket while reducing computation between content-uncorrelated elements.This design is identified as the source of its cost-effectiveness advantage.
- Clustering Scope: Local clustering performs better than global clustering because broader clustering can group less content-related elements and weaken subsequent attention.The reported explanation links global clustering to more complex, potentially linearly inseparable patterns.
Clustering Scope.
The parameter analysis balances reconstruction quality against computation and finds that local clustering is preferable to global clustering. The selected settings are σ = 0.5, R = 2, M = 16, and λ = 2.
- Parameter Analysis: The final parameter setting is σ = 0.5, R = 2, M = 16, and λ = 2.These values are selected to achieve a better trade-off between performance and computational complexity.
- Clustering Scope: Local clustering outperforms global clustering because global hashing can group less content-related elements, weakening bucket-wise attention.The resulting degradation is attributed to less effective MSA and poorer HSI restoration.
- Parameter Analysis: Increasing sparsity ratio σ lowers computational cost, but performance degrades dramatically when σ exceeds 50%.The analysis therefore identifies a trade-off between retaining patches and reducing computation.
- Parameter Analysis: Increasing round number R from 1 to 6 improves reconstruction quality, while gains diminish when R ≥ 2.The final configuration uses R = 2.
6 Conclusion
CST embeds HSI sparsity into learning-based reconstruction through informative-region detection and correlated-token clustering. Experiments report stronger performance, lower computational cost, and more detailed real-world reconstructions than existing algorithms.
- 6 Conclusion: CST detects informative HSI regions with SASM before applying SAH-MSA to cluster spatially scattered, closely correlated tokens.This coarse-to-fine design integrates sparsity into the reconstruction pipeline.
- 6 Conclusion: Extensive quantitative and qualitative experiments show CST significantly outperforms other state-of-the-art methods while requiring cheaper computational costs.
- 6 Conclusion: CST produces more visually pleasing real-world HSI reconstructions with finer-grained details and structural contents than existing algorithms.