[Paper Review] Coarse-to-Fine Sparse Transformer for Hyperspectral Image Reconstruction

Yuanhao Cai, Jing Lin+6 more|arXiv (Cornell University)|Mar 9, 2022

Photoacoustic and Ultrasonic Imaging5 citations

TL;DR

This paper proposes a coarse-to-fine sparse Transformer (CST) for hyperspectral image (HSI) reconstruction that leverages spatial sparsity in HSI signals to improve efficiency and performance. By using a spectra-aware screening mechanism (SASM) to select informative patches and a spectra-aggregation hashing multi-head self-attention (SAH-MSA) module to cluster content-correlated pixels, CST achieves state-of-the-art reconstruction quality with significantly reduced computational cost compared to existing methods.

ABSTRACT

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

Motivation & Objective

• To address the inefficiency and performance limitations of existing CNN and Transformer-based methods in hyperspectral image (HSI) reconstruction.
• To exploit the spatial sparsity inherent in HSI signals, which are often dominated by uninformative dark regions.
• To reduce computational cost by avoiding attention computation on uninformative or content-uncorrelated tokens.
• To improve long-range dependency modeling and non-local self-similarity capture in HSI reconstruction through content-aware clustering of tokens.
• To achieve a better trade-off between reconstruction quality and computational complexity.

Proposed method

• Proposes a spectra-aware screening mechanism (SASM) to identify and select only informative, signal-rich patches from the HSI for downstream processing.
• Applies a coarse-to-fine learning scheme: first selecting patches via SASM, then clustering pixels within each patch based on spectral similarity using a hashing function.
• Introduces a spectra-aggregation hashing multi-head self-attention (SAH-MSA) module that computes self-attention only within content-correlated clusters (buckets), reducing redundant computation.
• Uses a learnable hash function to map spectral features into discrete buckets, enabling efficient clustering of spatially scattered but spectrally similar pixels.
• Employs a multi-round refinement strategy (R rounds) to iteratively improve feature representation and reconstruction quality.
• Combines a multi-scale loss with a spectral consistency loss to preserve fine details and structural fidelity in reconstructed HSIs.

Experimental results

Research questions

• RQ1 How can the inherent spatial sparsity of hyperspectral images be effectively exploited in deep learning-based reconstruction models?
• RQ2 Can content-aware patch selection and token clustering improve the efficiency and accuracy of self-attention mechanisms in HSI reconstruction?
• RQ3 What is the optimal trade-off between computational cost and reconstruction performance when applying sparse attention to HSI data?
• RQ4 Does local clustering of tokens based on spectral similarity outperform global clustering in preserving structural and spectral details?
• RQ5 How does the proposed method compare to state-of-the-art CNN and Transformer-based approaches in terms of PSNR, SSIM, and visual quality?

Key findings

• The proposed CST achieves a PSNR of 35.53 dB on the CAEX dataset, outperforming the previous SOTA method by 0.96 dB with 28.7% fewer FLOPS.
• The SAH-MSA module achieves a 2.96 dB improvement over the baseline in PSNR while requiring only 0.85M additional parameters and 18.20G FLOPS, demonstrating superior cost-effectiveness.
• Local clustering scope (within M×M patches) yields better performance than global clustering, as larger scopes degrade hash clustering quality and reduce content correlation within buckets.
• The optimal sparsity rate σ is 0.5, as higher values (e.g., >50%) lead to significant performance degradation due to loss of critical signal information.
• The round number R=2 provides the best balance between performance and computational cost, with diminishing returns beyond R=2.
• Visual results show that CST preserves more fine-grained textures and structural details compared to existing methods, especially in complex scenes with sparse spectral content.