Empowering Snapshot Compressive Imaging:
Spatial-Spectral State Space Model with Across-Scanning and Local Enhancement
https://arxiv.org/pdf/2408.00629
Introduction
1. 背景与动机
快照压缩成像(SCI):
SCI技术依赖于解码算法(如CNN或Transformer)从压缩测量中重建高光谱图像(HSI)。现有的基于CNN和Transformer的方法已经证明了其有效性,但各有其局限性。
CNN方法虽然在提取图像局部特征方面具有优势,但无法充分建模长距离依赖关系。
Transformer方法能够实现全局感受野,相比CNN有更好的性能,但面临计算成本高的问题,这是由于注意力机制的二次复杂度所致。
Mamba模型:
近期Mamba模型在一些视觉任务中表现出了比CNN和Transformer更好的性能,但这些模型并未充分利用HSI在空间和光谱维度上的局部相似性。
同时,SSM(状态空间模型)在处理众多光谱带时,其长序列建模能力可能为HSI重建提供优势,但这一潜力尚未被充分挖掘。
2. 问题与挑战
当前HSI重建算法的不足主要体现在以下两个方面:
无法兼顾全局和局部信息:传统方法在全局和局部信息提取方面难以平衡。
高计算成本:特别是Transformer类方法,其二次复杂度的注意力机制导致了高昂的计算成本。
3. 方法与贡献
本文提出了一种名为ASLE-SSM的Spatial-Spectral SSM,结合了跨扫描和局部增强,用于HSI的重建。该模型主要贡献如下:
新颖的光谱-空间状态空间模型:
ASLE-SSM能够同时从光谱和空间维度整合信息,以充分建模空间上下文和光谱特征。
跨扫描方法:
该方法集成了来自空间和光谱维度的局部特征,在一个扫描过程中进行多通道交互,促进重建过程。
局部扫描方法:
在空间SSM模块中引入局部扫描方法,以进一步增强局部特征提取。
实验效果显著:
实验结果表明,ASLE-SSM在模拟和真实数据集上的表现优于当前最先进的方法,同时具有最少的参数数量、最低的计算成本和最快的推理速度。
Method
1. Architecture of ASLE-SSM
整体架构:
ASLE-SSM采用U型网络结构,以空间-光谱SSM块(Spatial-Spectral SSM Block)为基本单元。网络由编码器(Encoder)、瓶颈层(Bottleneck)和解码器(Decoder)组成。
输入处理:
输入信号为向量化的测量值y ∈ R^{H(W+d(Nλ-1))}
,通过反向色散和位移操作恢复为初始输入信号X ∈ R^{H×W×Nλ}
。
将输入X0
与物理掩码M
进行拼接,并通过1×1
卷积层整合原始信号和物理掩码的信息。
使用3×3
卷积嵌入层将整合后的输入X0
转换为特征X'0 ∈ R^{H×W×C}
。
编码器和解码器:
特征X'0
通过下采样和上采样过程,在编码器和解码器中传递。
解码器输出后,使用3×3
卷积映射获取残差信息,并与整合后的输入相加,得到最终的重建高光谱图像(HSI)。
模型尺寸:
通过调整编码器、瓶颈层和解码器中空间-光谱SSM块的数量,实现ASLE-SSM模型的小型、中型和大型尺寸。
Figure 2: The flowchart of the proposed ASLE-SSM. The framework is arranged by adopting a U-shaped structure. (a) Spatial-Spectral SSM Block. This module mainly consists of a spatial SSM and a spectral SSM, depicted in the sub-flowchart (b) and (c), respectively. (d) The specific method of LE-SSM in the Spatial SSM Module. Scanning directions are divided into global and local branches. (e) The illustration of our Across-Scanning Mechanism. The scanning process is demonstrated by labeling the order of each data point.
2. Spatial-Spectral SSM Block
空间SSM(Spatial SSM):
结合全局和局部扫描方法,通过行扫描同时进行全局和局部建模。
引入局部扫描(LE-SSM),在空间维度上增强局部特征提取,如图2(d)所示。
光谱-空间交叉SSM(Spectral-Spatial Cross SSM):
引入跨扫描机制(Across-Scanning),利用空间-光谱局部立方体来增强模块捕捉局部细节的能力,如图2(e)所示。
跨扫描机制通过在一个扫描过程中整合空间和光谱维度的局部特征,促进跨通道交互。
3. Across-Scanning Mechanism
跨扫描方法:
基于空间-光谱局部立方体,跨扫描机制通过整合相邻光谱带和像素之间的局部相似性来指导重建过程。
如图2(e)所示,跨扫描过程通过标记每个数据点的顺序来演示,实现空间和光谱维度的联合扫描。
4. Local Enhancement Mechanism
局部增强机制:
在空间SSM模块中,通过全局和局部扫描分支的结合,增强局部特征提取。
局部扫描分支专注于局部特征,而全局扫描分支则提供全局上下文信息。