[论文解读] Mask-guided Spectral-wise Transformer for Efficient Hyperspectral Image Reconstruction

Yuanhao Cai, Jing Lin+6 more|arXiv (Cornell University)|Nov 15, 2021

Photoacoustic and Ultrasonic Imaging参考文献 55被引用 15

一句话总结

本文提出了一种基于Transformer的新型框架——掩码引导光谱注意力变换器（MST），用于编码孔径快照光谱成像（CASSI）中高效高光谱图像（HSI）重建。通过在光谱维度上利用光谱注意力多头自注意力（S-MSA）建模自注意力机制，并借助掩码引导机制（MM）引导注意力，MST在显著降低计算和内存成本的同时，实现了最先进的性能表现。

ABSTRACT

Hyperspectral image (HSI) reconstruction aims to recover the 3D spatial-spectral signal from a 2D measurement in the coded aperture snapshot spectral imaging (CASSI) system. The HSI representations are highly similar and correlated across the spectral dimension. Modeling the inter-spectra interactions is beneficial for HSI reconstruction. However, existing CNN-based methods show limitations in capturing spectral-wise similarity and long-range dependencies. Besides, the HSI information is modulated by a coded aperture (physical mask) in CASSI. Nonetheless, current algorithms have not fully explored the guidance effect of the mask for HSI restoration. In this paper, we propose a novel framework, Mask-guided Spectral-wise Transformer (MST), for HSI reconstruction. Specifically, we present a Spectral-wise Multi-head Self-Attention (S-MSA) that treats each spectral feature as a token and calculates self-attention along the spectral dimension. In addition, we customize a Mask-guided Mechanism (MM) that directs S-MSA to pay attention to spatial regions with high-fidelity spectral representations. Extensive experiments show that our MST significantly outperforms state-of-the-art (SOTA) methods on simulation and real HSI datasets while requiring dramatically cheaper computational and memory costs. Code and pre-trained models are available at https://github.com/caiyuanhao1998/MST/

研究动机与目标

• 为解决基于CNN的方法在建模高光谱图像（HSI）重建中光谱相关性与长距离依赖关系方面的局限性。
• 更有效地利用CASSI系统中的物理编码孔径掩码，因为现有方法未能充分发挥其引导潜力。
• 在保持或提升重建质量的同时，降低HSI重建的计算与内存开销。
• 开发一种能够利用HSI数据固有的高光谱相关性与空间稀疏性的方法，实现高效建模。

提出的方法

• 提出光谱注意力多头自注意力（S-MSA），将每个光谱波段视为一个token，在光谱维度上计算自注意力，以建模波段间的相关性。
• 提出掩码引导机制（MM），利用编码孔径掩码引导S-MSA聚焦于具有更高保真度光谱表示的空间区域。
• 用光谱注意力替代标准的空间注意力机制，以更好地匹配HSI数据内在的光谱相似性。
• 采用改进的输入方案，掩码不通过逐元素相乘直接应用于输入，避免对HSI信息造成破坏。
• 通过聚焦于光谱依赖关系而非空间长程交互，构建轻量化架构，显著降低FLOPS与参数量。
• 将S-MSA与MM集成至基于Transformer的编码器-解码器框架中，实现端到端的HSI重建。

实验结果

研究问题

• RQ1 在HSI重建中，沿光谱维度建模自注意力是否能比空间注意力更好地捕捉波段间相关性与长距离依赖？
• RQ2 物理编码孔径掩码在多大程度上能有效引导注意力机制，以提升重建质量与效率？
• RQ3 与标准Transformer相比，光谱注意力机制是否能在保持或提升重建性能的同时降低计算成本？
• RQ4 所提出的掩码引导机制（MM）是否能优于传统输入方案（即通过掩码相乘破坏HSI数据）？
• RQ5 与全局或局部窗口基Multi-Head Self-Attention（MSA）相比，S-MSA模块在光谱相关性建模方面提升程度如何？

主要发现

• 所提出的MST框架在CAESR数据集上达到33.17 dB的PSNR，优于先前SOTA方法0.89 dB。
• 仅使用0.70M参数与2.93G FLOPS的S-MSA模块，相比基线模型实现0.89 dB的性能提升，且计算成本远低于全局MSA。
• 使用S-MSA重建的HSI相关系数图与真实值最接近，证实其在建模光谱相似性方面的有效性。
• 仅使用掩码引导机制（MM）即可使PSNR提升1.08 dB，表明其能有效引导注意力聚焦于高保真区域并抑制噪声。
• 结合S-MSA与MM后，相比基线模型实现1.08 dB的PSNR提升，且在PSNR与效率指标上全面超越所有SOTA方法。
• 视觉分析表明，MM模块增强了对结构细节与纹理的关注，有效减少了重建HSI中的模糊与伪影。