[论文解读] S^2-Transformer for Mask-Aware Hyperspectral Image Reconstruction

Jiamian Wang, Kunpeng Li+3 more|arXiv (Cornell University)|Sep 24, 2022

Photoacoustic and Ultrasonic Imaging被引用 4

一句话总结

本文提出一种具有掩码感知学习的S²-Transformer网络，用于从CASSI测量中实现高保真度的高光谱图像重建。通过联合建模空间与光谱注意力，并自适应地优先处理掩码区域的重建损失，该方法实现了最先进性能，相较于先前方法在PSNR上提升0.17 dB，在SSIM上提升0.0015。

ABSTRACT

Snapshot compressive imaging (SCI) surges as a novel way of capturing hyperspectral images. It operates an optical encoder to compress the 3D data into a 2D measurement and adopts a software decoder for the signal reconstruction. Recently, a representative SCI set-up of coded aperture snapshot compressive imager (CASSI) with Transformer reconstruction backend remarks high-fidelity sensing performance. However, dominant spatial and spectral attention designs show limitations in hyperspectral modeling. The spatial attention values describe the inter-pixel correlation but overlook the across-spectra variation within each pixel. The spectral attention size is unscalable to the token spatial size and thus bottlenecks information allocation. Besides, CASSI entangles the spatial and spectral information into a 2D measurement, placing a barrier for information disentanglement and modeling. In addition, CASSI blocks the light with a physical binary mask, yielding the masked data loss. To tackle above challenges, we propose a spatial-spectral (S2-) Transformer implemented by a paralleled attention design and a mask-aware learning strategy. Firstly, we systematically explore pros and cons of different spatial (-spectral) attention designs, based on which we find performing both attentions in parallel well disentangles and models the blended information. Secondly, the masked pixels induce higher prediction difficulty and should be treated differently from unmasked ones. We adaptively prioritize the loss penalty attributing to the mask structure by referring to the mask-encoded prediction as an uncertainty estimator. We theoretically discuss the distinct convergence tendencies between masked/unmasked regions of the proposed learning strategy. Extensive experiments demonstrate that on average, the results of the proposed method are superior over the state-of-the-art method.

研究动机与目标

• 解决基于CASSI的高光谱成像中两个主要的数据损失问题：光谱通道混叠以及由于物理编码孔径导致的像素丢失。
• 通过新颖的Transformer架构，利用高光谱数据中固有的空间与光谱相关性，提升重建保真度。
• 建模由光学掩蔽引起的像素级重建难度，将其与由纹理引起的不确定性区分开来。
• 通过基于物理成像原理的设计，增强深度学习模型在高光谱图像重建中的可解释性。

提出的方法

• 提出一种S²-Transformer，采用并行的空间与光谱自注意力机制，以解耦2D CASSI测量中纠缠的光谱与空间信息。
• 引入掩码编码（ME）损失项，基于物理掩码图案对网络早期层的预测进行惩罚。
• 设计掩码感知（MA）损失，通过从预测输出动态推断像素级重建难度，并对损失权重更高的掩码区域进行优先处理。
• 采用双分支架构，使空间与光谱注意力并行处理，并通过特征拼接与残差连接实现交互。
• 在预训练阶段使用ME损失，以稳定早期训练并提升掩码区域的收敛性。
• 理论分析表明，在所提出的损失策略下，掩码区域与未掩码区域表现出不同的收敛行为。

实验结果

研究问题

• RQ1 如何联合设计空间与光谱注意力机制，以更好地利用2D压缩测量中高光谱数据的特性？
• RQ2 与标准L1/L2损失相比，建模由光学掩蔽引起的不确定性的程度在多大程度上能提升重建保真度？
• RQ3 是否一种能够根据预测置信度自适应调整的动态掩码感知损失策略，可提升掩码区域的收敛性与性能？
• RQ4 空间与光谱注意力之间的交互如何影响模型的表征能力与可解释性？

主要发现

• 所提出的S²-Transformer结合掩码感知学习在基准数据集上实现了36.48 dB的PSNR与0.9584的SSIM，相较于先前最先进方法，PSNR提升0.17 dB，SSIM提升0.0015。
• 消融实验表明，同时使用ME与MA损失可获得最佳性能，而单独使用任一损失均仅带来微小或无改善。
• 掩码感知（MA）损失显著降低了掩码区域的绝对重建误差，如收敛曲线与视觉对比所示。
• 视觉结果表明，掩码感知学习显著提升了纹理恢复与细节重建能力，尤其在模拟与真实世界数据的掩码区域表现更优。
• 理论与实证分析证实，所提出的损失策略使掩码区域的收敛速度更快、更稳定，相较于未掩码区域表现更优。
• 并行的空间-光谱注意力机制有效捕捉了长程光谱依赖性与空间相关性，相关性矩阵分析与注意力可视化结果均支持此结论。