PaperlayerView exact revision on arXiv

Source-linked AI summary

MST++: Multi-stage Spectral-wise Transformer for Efficient Spectral Reconstruction

Yuanhao Cai, Jing Lin, Zudi Lin, Haoqian Wang, Yulun Zhang, Hanspeter Pfister, Radu Timofte, Luc Van Gool

arXiv:2204.07908v1cs.CV

TL;DR

CNN-based spectral reconstruction performs well but is limited in capturing long-range dependencies and inter-spectral self-similarity. MST++ addresses this with spectral-wise Transformer attention and cascaded multi-stage processing, achieving strong efficiency and challenge performance. The paper reports that MST++ significantly outperforms state-of-the-art methods and wins first place in the NTIRE 2022 Spectral Reconstruction Challenge.

  • Problem

    CNN-based spectral reconstruction methods have limited ability to capture long-range dependencies and inter-spectral self-similarity.

  • Method

    MST++ uses S-MSA within SABs and U-shaped SSTs, then cascades several SSTs for coarse-to-fine spectral reconstruction from RGB images.

  • Results

    MST++ significantly outperforms state-of-the-art methods with cheaper Params and FLOPS and wins first place in the NTIRE 2022 Spectral Reconstruction Challenge.

  • Takeaways & Limitations

    MST++ provides an efficient Transformer-based framework for spectral reconstruction and is released with code and pretrained models for further research.

Abstract

from arXiv · show

Existing leading methods for spectral reconstruction (SR) focus on designing deeper or wider convolutional neural networks (CNNs) to learn the end-to-end mapping from the RGB image to its hyperspectral image (HSI). These CNN-based methods achieve impressive restoration performance while showing limitations in capturing the long-range dependencies and self-similarity prior. To cope with this problem, we propose a novel Transformer-based method, Multi-stage Spectral-wise Transformer (MST++), for efficient spectral reconstruction. In particular, we employ Spectral-wise Multi-head Self-attention (S-MSA) that is based on the HSI spatially sparse while spectrally self-similar nature to compose the basic unit, Spectral-wise Attention Block (SAB). Then SABs build up Single-stage Spectral-wise Transformer (SST) that exploits a U-shaped structure to extract multi-resolution contextual information. Finally, our MST++, cascaded by several SSTs, progressively improves the reconstruction quality from coarse to fine. Comprehensive experiments show that our MST++ significantly outperforms other state-of-the-art methods. In the NTIRE 2022 Spectral Reconstruction Challenge, our approach won the First place. Code and pre-trained models are publicly available at https://github.com/caiyuanhao1998/MST-plus-plus.

1. Introduction

Spectral reconstruction from RGB images addresses the cost and speed limitations of hyperspectral imaging, but CNN-based methods struggle with long-range dependencies and inter-spectral self-similarity. MST++ introduces a Transformer-based, multi-stage solution and reports strong performance with lower computational and memory costs.

  • Motivation: Spectrometer-based HSI acquisition is time-consuming and limits use in dynamic or real-time scenes.Scanning along spatial or spectral dimensions impedes broader application.
  • Motivation: Spectral reconstruction reconstructs HSI cubes from RGB images, which can be obtained using relatively inexpensive RGB cameras.This avoids the expensive hardware required by snapshot compressive imaging systems.
  • Limitations of Existing Methods: CNN-based SR methods achieve impressive performance but have limited ability to capture long-range dependencies and inter-spectra self-similarity.Earlier model-based approaches also suffer from limited representation capacity and poor generalization.
  • Limitations of Existing Methods: Standard Transformers model spatial interactions, although global attention is quadratic in spatial size and local attention has position-specific receptive fields.HSIs are spatially sparse but spectrally self-similar, motivating attention along the spectral dimension.
  • Proposed Framework: MST++ uses S-MSA-based SABs, U-shaped SSTs, and cascaded stages to progressively improve reconstruction from coarse to fine.S-MSA treats each spectral feature map as a token, while SSTs extract multi-resolution spectral context.
  • Results and Resources: MST++ dramatically outperforms state-of-the-art methods while requiring much cheaper Params and FLOPS, and wins first place in the NTIRE 2022 challenge.The authors also report publicly available code and pretrained models.

2. Related Work

Hyperspectral imaging offers richer spectral detail but conventional scanning is slow and physically cumbersome. Related spectral reconstruction work progresses from handcrafted priors and CNNs toward Transformers, while CNN methods remain limited in modeling non-local dependencies and self-similarity.

  • Hyperspectral Image Acquisition: Traditional HSI scanners require long acquisition times and are physically too large for portable platforms.These constraints make them unsuitable for measuring dynamic scenes.
  • Spectral Reconstruction from RGB: Conventional SR methods use handcrafted hyperspectral priors, sparse coding, or shallow learning models but have limited representation capacity and poor generalization.Examples include data sparsity expansion, dictionary-based sparse coding, and spectral-prior learning models.
  • Spectral Reconstruction from RGB: CNN-based SR methods learn RGB-to-HSI mappings and achieve impressive results but struggle with non-local self-similarity and long-range inter-dependencies.Representative approaches include HSCNN, HSCNN-R, and pixel-aware function-mixture networks.
  • Transformer-Based Vision Models: Transformer methods are attractive because self-attention captures long-range correlations between spatial regions and has been applied across vision tasks.The related work positions Transformers as a growing approach in both high-level and low-level vision.

3. Method

MST++ maps an RGB image to an HSI using cascaded U-shaped SSTs. Its SAB units combine feed-forward processing, normalization, and spectral-wise self-attention tailored to spectral correlations.

  • Network Architecture: MST++ cascades multiple SSTs to reconstruct an HSI from an RGB input, using a long identity mapping to ease training.Each SST contains an encoder, bottleneck, and decoder in a U-shaped architecture.
  • Network Architecture: SSTs use downsampling, SABs, a bottleneck, symmetric upsampling, and skip connections to extract multi-resolution spectral context.Skip connections reduce information loss from downsampling.
  • Spectral-wise Attention Block: Each SAB combines an FFN, S-MSA, and two layer-normalization layers.The S-MSA details define the spectral attention mechanism used inside the block.
  • Spectral-wise Multi-head Self-Attention: S-MSA reshapes input feature maps into tokens and linearly projects them into query, key, and value representations.The projections use learnable matrices WQ, WK, and WV.
  • Spectral-wise Multi-head Self-Attention: S-MSA splits query, key, and value into heads along the spectral channel dimension and computes attention for each head.Each spectral representation is treated as a token rather than a spatial region.
  • Spectral-wise Multi-head Self-Attention: A learnable σj adapts attention weights because spectral density varies substantially across wavelengths.This parameter reweights the matrix multiplication inside each attention head.

S - MSA}(\

After head-wise spectral attention, MST++ concatenates the head outputs, applies a linear projection, adds position embedding, and reshapes the result into output feature maps.

  • S-MSA: The outputs of all attention heads are concatenated, linearly projected, and combined with a position embedding.The embedding uses depth-wise convolutions and GELU processing to encode spectral-channel position.
  • S-MSA: The aggregated representation is reshaped to produce output feature maps Xout with spatial and spectral dimensions.The position embedding encodes channel order because HSI channels are sorted by wavelength.

3.3. Discussion with Original Transformers

The discussion contrasts spatial-wise and spectral-wise self-attention, motivating S-MSA for HSI reconstruction because it models spectral correlations with global receptive fields and linear spatial cost.

  • General Paradigm of MSA: The general MSA pipeline projects input X into query, key, and value matrices, splits them into heads, applies attention, and linearly projects the concatenated outputs.The formulation omits positional embeddings when comparing the main distinction between spatial-wise MSA and S-MSA.
  • Global MSA: Global MSA models all spatial tokens, giving a quadratic cost of 2(HW)^2 C and a large receptive field.Its redundant key sampling can cause over-smooth results and nonconvergence issues.
  • Window-based MSA: W-MSA reduces complexity to 2M^2HWC by restricting self-attention to position-specific windows.This linear spatial cost comes with limited receptive fields that may miss related non-local tokens.
  • Motivation: HSI representations are spatially sparse but spectrally similar and correlated, making spectral-wise interactions more cost-effective than spatial-wise interactions.This characteristic motivates adopting S-MSA for spectral reconstruction.
  • S-MSA: S-MSA treats each spectral feature map as a token and computes self-attention along the spectral dimension, adapting attention with learnable σj values.The learnable parameter reweights K_j^T Q_j to account for wavelength-dependent spectral density.
  • S-MSA: S-MSA has complexity 2HWC^2/N, linear in spatial size, while retaining global receptive fields through globally represented feature maps.It models spectral self-similarity and is presented as more cost-effective than global and window-based MSA.

3.4. Ensemble Strategy

MST++ uses three ensemble strategies in the NTIRE 2022 challenge: transformed-input self-ensemble, multi-scale training, and Top-K multi-model fusion.

  • Self-ensemble: Self-ensemble averages outputs from inputs flipped or rotated through multiple spatial transformations.Each transformed output is returned to the original orientation before averaging.
  • Multi-scale ensemble: Multi-scale ensemble averages whole-image outputs from models trained with 256×256, 128×128, and 64×64 patches.The three patch sizes provide the training scales used for restoration-quality improvement.
  • Top-K multi-model ensemble: Top-K multi-model ensemble fuses reconstructions from the strongest models among MIRNet, MPRNet, Restormer, HINet, and MST families.The selected models are combined using weighted reconstructed HSIs.

4. Experiment

The experiments use 1000 RGB-HSI pairs with 31 wavelengths, synthetic shot noise, and standardized training and evaluation procedures. MRAE is the training objective because it is the deciding challenge metric.

  • Dataset: The NTIRE 2022 dataset contains 1000 RGB-HSI pairs, split 18:1:1 into train, validation, and test subsets, with 31 wavelengths from 400 nm to 700 nm.
  • Dataset: Generated RGB images receive shot noise to simulate real-camera conditions.
  • Implementation Details: Training uses 128 × 128 crops, batch size 20, Adam with β1 = 0.9 and β2 = 0.999, a 0.0004 initial learning rate, cosine annealing, and 300 epochs.
  • Implementation Details: MST++ reconstructs one 482×512×3 RGB image in 102.48 ms on an RTX 3090 GPU during testing.
  • Evaluation Metrics: MRAE measures pixel-wise disparity across wavelengths, while RMSE measures root mean squared reconstruction error.
  • Evaluation Metrics: MRAE is used as the training objective because it is the deciding metric for the NTIRE 2022 Spectral Reconstruction Challenge.

4.3. Main Results

MST++ is evaluated against spectral reconstruction, snapshot compressive imaging, and natural-image restoration baselines on validation and test settings. It combines strong reconstruction quality with low computational cost and wins the NTIRE 2022 challenge.

  • Quantitative Results on Valid Set: MST++ significantly outperforms nine SOTA baselines on the validation set while requiring the fewest Params and FLOPS.
  • Quantitative Results on Valid Set: 3.10 dB PSNR improvement over AWAN requires 40.10% of its Params and 8.52% of its FLOPS.
  • Quantitative Results on Valid Set: 7.43 dB PSNR improvement over HRNet requires 5.11% of its Params and 14.07% of its FLOPS.
  • Quantitative Results on Valid Set: 7.96 dB PSNR improvement over HSCNN+ requires 34.84% of its Params and 7.57% of its FLOPS.
  • Efficiency Comparison: MST++ occupies the top-left of the PSNR-Params-FLOPS comparison, indicating high performance with low computational and memory costs.
  • Test-Set Results: MST++ won the NTIRE 2022 Spectral Reconstruction Challenge championship among 231 participants.
  • Qualitative Results: Qualitative comparisons show sharper edges, preserved homogeneous-region smoothness, and spectral curves with the highest ground-truth correlation and coincidence.

4.4. Ablation Study

Ablations show that spectral-wise self-attention provides the strongest accuracy improvement with the lowest stated memory and computational costs. Three stages are selected, and ensemble strategies further improve MRAE.

  • Self-Attention Mechanism: The ablation baseline is formed by removing S-MSA from MST++ on the validation subset.
  • Self-Attention Mechanism: S-MSA yields the largest MRAE improvement, 0.1532, among SW-MSA, W-MSA, G-MSA, and S-MSA alternatives.
  • Self-Attention Mechanism: S-MSA requires the least memory and computational cost, while the alternatives increase FLOPS by 6.42, 6.42, and 7.43 GFLOPS for SW-MSA, W-MSA, and G-MSA.
  • Self-Attention Mechanism: The results are attributed to HSI spatial sparsity and spectral self-similarity, making inter-spectra dependency modeling more cost-effective than spatial-region correlation modeling.
  • Stage Number: MST++ achieves peak performance with Ns = 3, so the final spectral reconstruction model uses three stages.
  • Ensemble Strategies: Self-ensemble, multi-scale ensemble, and Top-K multi-model ensemble improve MRAE by 0.015, 0.033, and 0.045, respectively.

5. Future Work

MST++ addresses the lack of a low-cost, high-accuracy open-source SR baseline by releasing source code and pre-trained models. The authors aim to provide a model zoo and toolbox for the research community.

  • The authors identify the absence of a low-cost, high-accuracy open-source baseline for spectral reconstruction research.
  • Source code and pre-trained models for 11 state-of-the-art methods are publicly available.
  • The released resources are intended to serve as a model zoo and toolbox for the research community.

6. Conclusion

MST++ is a multi-stage Transformer framework for spectral reconstruction from RGB that progressively refines reconstruction from coarse to fine. Experiments report substantial advantages over state-of-the-art methods, with lower memory and computational costs, and first place in the NTIRE 2022 challenge.

  • MST++ cascades several SSTs to progressively improve spectral reconstruction quality from coarse to fine.
  • S-MSA treats each spectral feature map as a token, forming SABs that build the SST stages.
  • MST++ dramatically surpasses state-of-the-art methods while requiring cheaper memory and computational costs.
  • MST++ won first place in the NTIRE 2022 Challenge on Spectral Reconstruction from RGB.
Loading 2204.07908v1…