[论文解读] Coarse-to-Fine Sparse Transformer for Hyperspectral Image Reconstruction

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

Photoacoustic and Ultrasonic Imaging被引用 5

一句话总结

本文提出一种粗到细的稀疏Transformer（CST），用于高光谱图像（HSI）重建，通过利用HSI信号中的空间稀疏性，提升效率与性能。通过使用光谱感知筛选机制（SASM）选择信息丰富的图块，以及通过光谱聚合哈希多头自注意力（SAH-MSA）模块对内容相关像素进行聚类，CST在显著降低计算成本的同时，实现了最先进的重建质量。

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

研究动机与目标

• 解决现有基于CNN和Transformer的方法在高光谱图像（HSI）重建中的效率低下与性能受限问题。
• 利用HSI信号中固有的空间稀疏性，这些信号通常由大量无信息的暗区主导。
• 通过避免对无信息或内容无关的标记进行注意力计算，降低计算成本。
• 通过基于内容感知的标记聚类，提升长距离依赖建模与非局部自相似性捕捉能力，从而改进HSI重建。
• 在重建质量与计算复杂度之间实现更优的权衡。

提出的方法

• 提出一种光谱感知筛选机制（SASM），用于识别并仅选择HSI中信息丰富、信号强的图块用于后续处理。
• 采用粗到细的学习策略：首先通过SASM选择图块，然后利用哈希函数基于光谱相似性对每个图块内的像素进行聚类。
• 引入一种光谱聚合哈希多头自注意力（SAH-MSA）模块，仅在内容相关聚类（桶）内计算自注意力，从而减少冗余计算。
• 使用可学习的哈希函数将光谱特征映射到离散桶中，实现对空间分散但光谱相似像素的高效聚类。
• 采用多轮优化策略（R轮），迭代提升特征表示与重建质量。
• 结合多尺度损失与光谱一致性损失，以保留重建HSI中的精细细节与结构保真度。

实验结果

研究问题

• RQ1 如何在基于深度学习的重建模型中有效利用高光谱图像的固有空间稀疏性？
• RQ2 基于内容感知的图块选择与标记聚类是否能提升HSI重建中自注意力机制的效率与准确性？
• RQ3 在将稀疏注意力应用于HSI数据时，计算成本与重建性能之间的最优权衡是什么？
• RQ4 基于光谱相似性的局部标记聚类是否优于全局聚类，以更好地保留结构与光谱细节？
• RQ5 与当前最先进的基于CNN和Transformer的方法相比，所提方法在PSNR、SSIM和视觉质量方面表现如何？

主要发现

• 所提出的CST在CAEX数据集上达到35.53 dB的PSNR，优于先前SOTA方法0.96 dB，且FLOPS减少28.7%。
• SAH-MSA模块在PSNR上相比基线提升2.96 dB，仅增加0.85M参数与18.20G FLOPS，展现出卓越的成本效益。
• 局部聚类范围（在M×M图块内）优于全局聚类，因为更大的范围会降低哈希聚类质量，并减少桶内内容相关性。
• 最优稀疏率σ为0.5，因为更高的值（如>50%）会导致关键信号信息丢失，从而引起性能显著下降。
• 当R=2时，性能与计算成本达到最佳平衡，超过R=2后收益递减。
• 视觉结果表明，与现有方法相比，CST能更好地保留细粒度纹理与结构细节，尤其在光谱内容稀疏的复杂场景中。