MST++: Multi-stage Spectral-wise Transformer for Efficient Spectral Reconstruction (2204.07908v1)

Abstract: 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.

• The paper introduces spectral-wise self-attention that treats each spectral band as a token to enhance feature extraction in hyperspectral reconstruction.
• It employs a multi-stage Transformer design that progressively refines image quality while significantly reducing computational resources.
• Experimental results on the NTIRE 2022 dataset show improved PSNR and MRAE, establishing MST++ as a state-of-the-art method.

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

The paper presents MST++, a Transformer-based framework developed for enhancing spectral reconstruction (SR) performance. SR is the process of reconstructing hyperspectral images (HSI) from conventional RGB images. This capability is essential due to the broad applications of HSIs in fields like medical imaging and remote sensing. The proposed methodology addresses the limitations of prior convolutional neural networks (CNNs) that have predominantly been used for SR tasks.

Key Contributions

1. Spectral-wise Multi-head Self-attention (S-MSA): MST++ introduces S-MSA, which leverages the spatial sparsity and spectral self-similarity inherent in HSIs. Unlike typical Transformer architectures that focus on spatial dependencies, S-MSA treats each spectral band as a token, computing self-attention across spectral dimensions.
2. Spectral-wise Attention Block (SAB): These blocks form the core computational unit of the MST++ framework. SABs capture spectral dependencies efficiently, providing more targeted attention mechanism suitable for HSI characteristics.
3. Multi-stage Structure: The architecture incorporates several Single-stage Spectral-wise Transformers (SSTs) arranged in a multi-stage manner. This iterative strategy allows the model to progressively refine the spectral reconstruction from coarse to fine levels, enhancing the quality of the output HSIs.
4. Computational Efficiency: The framework significantly reduces the computational burden (lower FLOPS and Params) compared to existing methods, as evidenced by superior performance metrics like PSNR and MRAE with reduced computational resources.

Experimental Results

The MST++ framework underwent rigorous testing against state-of-the-art methods using the NTIRE 2022 Spectral Reconstruction Challenge dataset. It registered notable improvements in reconstruction accuracy, achieving first place in the competition. This performance validates the model's effectiveness in handling spectral reconstruction challenges with high computational efficiency.

Implications and Future Directions

The introduction of Transformer-based models like MST++ in SR tasks highlights a paradigm shift from traditional CNN-based approaches. The capability of effectively modeling long-range dependencies in the spectral domain is particularly advantageous.

Practical Implications:

• The deployment of MST++ can lead to more accurate and efficient HSI reconstruction in real-time applications, benefiting fields requiring prompt spectral data analysis.
• Its efficiency in computational resource usage aligns well with applications in environments with limited computational infrastructure.

Theoretical Implications:

• This research demonstrates the potential of adopting attention mechanisms focused on inter-spectral relationships in applications that involve multi-dimensional data structures.
• It encourages the pursuit of further Transformer adaptations that cater to domain-specific characteristics in visual data processing tasks.

Speculation on Future Developments:

• Future research may extend MST++ by integrating more sophisticated attention mechanisms or exploring hybrid models that blend the strengths of both CNNs and Transformers.
• Additionally, exploring domain adaptation techniques to generalize MST++ across various spectral imaging tasks beyond those discussed could be a promising avenue.

In conclusion, this work underscores the transformative potential of Transformers in spectral reconstruction and paves the way for continued exploration of attention mechanisms in similar analytical domains.