HAMSA：基于谱自适应的动态视觉状态空间模型原理与实践

1. 从序列建模到视觉理解：为什么我们需要新的视觉骨干网络？

最近在刷论文和开源社区的时候，一个词出现的频率越来越高：状态空间模型。如果你关注过自然语言处理，尤其是长序列建模，那你对Mamba这个名字一定不陌生。它凭借其选择性状态空间模型，在处理长文本时展现出了媲美甚至超越Transformer的效率与性能。但视觉领域呢？我们还在被Transformer和CNN统治着。卷积神经网络擅长局部特征提取，但建模长距离依赖的能力有限；视觉Transformer通过自注意力机制实现了全局交互，但计算复杂度随着图像分辨率呈平方级增长，对高分辨率图像（如医学影像、遥感图像）的处理成本高得吓人。

这就引出了一个核心问题：能否将SSM在序列建模上的高效优势，平移到视觉任务中？这不仅仅是“拿来就用”那么简单。图像数据是二维的，具有强烈的空间局部性和结构性，而语言序列是一维的。直接将一维SSM应用到图像块序列上，会忽略掉像素之间固有的二维邻接关系，效果往往不尽人意。因此，过去一年里，涌现了像Vision Mamba、VMamba等一系列工作，它们尝试通过扫描策略（如双向扫描）来模拟二维空间关系，取得了一些进展。

然而，这些方法大多将SSM视为一个“黑盒”算子，其核心参数（如状态矩阵A）是静态的或通过简单网络生成的。这里存在一个根本性的矛盾：视觉内容具有高度的动态性和多样性。一张图片里，天空是平滑的，纹理是高频的，边缘是尖锐的。一个固定的、或由简单线性层生成的SSM参数，真的能最优地建模所有这些截然不同的视觉模式吗？答案很可能是否定的。这就好比用一把固定焦距的镜头去拍摄所有场景，远景和微距都无法兼顾。

HAMSA这篇工作，正是瞄准了这个痛点。它的全称是“Hierarchical Adaptive Mamba with Spectral-Awareness”，直译过来是“具有谱感知能力的层次化自适应Mamba”。这个名字已经透露了它的两大核心创新点：“谱自适应”和“层次化”。简单来说，它不再使用“一刀切”的SSM参数，而是让模型能够根据输入图像内容本身，动态地、自适应地调整SSM的核心参数，并且这种调整是建立在图像频域特性分析的基础上的。同时，它保持了经典的层次化金字塔架构，以适应多尺度视觉表征的需求。

我花了一些时间研读论文和代码，并尝试在图像分类和分割任务上复现。实话说，这种“内容感知”的动态机制带来的性能提升是显著的，尤其是在处理复杂纹理和细节丰富的图像时。下面，我就结合自己的理解，拆解一下HAMSA到底是怎么做的，以及我们在复现和应用时需要注意哪些关键细节。

2. HAMSA架构全景：层次化设计如何融入动态SSM

HAMSA的整体架构遵循了现代视觉骨干网络的设计范式，即层次化金字塔结构。这种结构的好处显而易见：浅层捕获细节和边缘，深层捕获语义和物体。对于下游任务如目标检测、语义分割，多尺度特征图是必不可少的。HAMSA采用了经典的四阶段设计，随着网络加深，空间分辨率逐渐降低，通道数逐渐增加。

2.1 从Patch Embedding到HAMSA Block

输入一张H×W×3的RGB图像，首先经过一个Patch Embedding层。这里和ViT类似，使用一个卷积核大小和步长均为4的卷积层，将图像分割成不重叠的块，并将每个块投影到C维嵌入空间。这样，我们得到了一个 (H/4) × (W/4) × C 的序列，可以将其视为长度为 L = (H/4)*(W/4) 的一维序列，但我们需要时刻记住其二维来源。

接下来，序列会依次通过四个阶段。每个阶段都由若干个相同的HAMSA Block堆叠而成，并且在阶段之间，会插入一个Patch Merging层来降低分辨率、增加通道数，实现下采样。

HAMSA Block是整个架构的核心单元。它采用了类似Swin Transformer的残差结构，但核心的注意力模块被替换为了我们自适应的SSM模块。一个标准的HAMSA Block可以表示为：

Z = X + SSM_Adaptive( LayerNorm(X) ) Y = Z + MLP( LayerNorm(Z) )

其中，SSM_Adaptive就是我们即将深入探讨的、具备谱自适应能力的状态空间模型层。

2.2 动态与静态的权衡：什么参数应该被自适应？

在深入自适应机制前，我们需要先回顾一下SSM（特别是离散化后的）的核心公式。对于一个输入序列x(t)，SSM通过一个隐状态h(t)将其映射到输出y(t)。其连续形式由参数(A, B, C)决定。经过零阶保持离散化后，我们得到：

h_t = \bar{A} * h_{t-1} + \bar{B} * x_t y_t = C * h_t

其中，\bar{A} = exp(Δ * A),\bar{B} = (Δ * A)^{-1} (exp(Δ * A) - I) * Δ * B。这里，Δ是一个重要的步长参数，控制着离散化的粒度。

在传统的Mamba或S4模型中，A矩阵通常是固定初始化的（如HiPPO初始化），或者B, C由线性层根据输入x_t动态生成。Δ则通常由另一个参数化网络根据x_t生成。HAMSA的思考更进一步：A矩阵，这个决定系统动态（即如何记忆和遗忘信息）的核心参数，是否也应该根据输入内容动态变化？

答案是肯定的。A矩阵的特征值决定了SSM的频率响应特性。一个能够根据图像局部频谱特性动态调整A的模型，理论上可以更灵活地选择“记住”哪些空间频率的信息，“过滤”掉哪些信息。这就是“谱自适应”概念的由来。

3. 谱自适应机制核心：从空域到频域的感知与调制

HAMSA实现“谱自适应”的关键，在于它设计了一个轻量级的谱感知头，用于生成动态的A矩阵参数。这个过程不是直接操作庞大的A矩阵（其维度为状态维度N），而是通过一个更紧凑的、与频域相关的参数化方式。

3.1 局部频谱特征的提取

首先，对于输入特征图X ∈ R^{H×W×C}，我们不是直接处理整个图，而是在局部窗口内进行操作，这与Swin Transformer的窗口注意力思想一脉相承，能大幅降低计算量并保留局部性。

假设我们在一个M×M的局部窗口内操作。对于窗口内的C个特征通道，我们将其视为M^2个空间位置。要分析其频谱特性，一个直接的想法是使用二维离散傅里叶变换。然而，对每个窗口、每个样本都做FFT计算开销较大，并且我们不需要完整的频域表示，只需要一个能概括其主要频谱特性的描述符。

因此，HAMSA采用了一种更工程化的方法：使用一组可学习的空间滤波器（卷积核）来模拟对不同频率带的响应。具体来说，网络学习K个小的卷积核（例如3x3）。将这些核分别与窗口内的特征图进行卷积操作，然后对输出进行空间池化（如全局平均池化），得到K个标量响应值。这K个值构成了一个频谱描述向量f ∈ R^K。你可以把这K个核理解为针对不同方向、不同尺度（可粗略对应频率）的边缘或纹理检测器，它们的响应强弱共同刻画了该局部窗口的频谱能量分布。

注意：这里的设计非常巧妙。它避免了显式的FFT计算，将频域感知能力嵌入到了可学习的卷积核中，使得整个过程完全可微分，能够端到端训练。K是一个超参数，通常设置为一个较小的数（如8或16），以保证效率。

3.2 从频谱描述到动态A矩阵

得到频谱描述向量f后，我们需要将其转化为对SSMA矩阵的调整量。直接生成一个N×N的矩阵是不现实的。HAMSA利用了SSM中常用的一种结构化的A矩阵形式——对角加低秩，或者直接使用对角矩阵。这是因为在S4/Mamba中，为了计算效率，A通常被约束为对角矩阵（或对角加上下三角矩阵），其离散化\bar{A}的计算可以逐元素进行。

假设我们使用对角矩阵A = diag(λ_1, λ_2, ..., λ_N)，其中λ_i是复数，通常位于左半复平面（保证系统稳定）。那么，自适应就变成了：如何根据频谱描述f，生成一组调整量Δλ_i，来微调每个通道的λ_i？

HAMSA通过一个小的MLP（多层感知机）来实现：

Δλ = MLP_Spectrum(f)

这个MLP以f为输入，输出一个N维的复数向量Δλ ∈ C^N。然后，我们将这个调整量加到基础（或共享）的A矩阵参数上：

λ_adapted = λ_base + Δλ

这里，λ_base是所有HAMSA Block共享的、可学习的基础参数，初始值可以采用HiPPO等经典方法。而Δλ则是每个窗口、每个样本、每个网络层都不同的动态调整部分。

3.3 自适应过程的计算图与梯度流

理解这个过程的梯度流很重要。λ_base是全局可学习参数，其梯度来自于所有样本和所有位置的反向传播。而Δλ是通过一个由频谱描述向量f参数化的网络生成的，f又依赖于当前输入特征X和那K个可学习的频谱感知卷积核。

因此，整个系统形成了一个闭环：

1. 输入特征X决定了局部频谱特性f。
2. f通过MLP决定了SSM动态系统参数A的调整量Δλ。
3. 调整后的SSM (A_adapted) 处理特征X，产生输出。
4. 输出的损失梯度会同时更新：特征提取权重、频谱感知卷积核、MLP_Spectrum的权重、以及基础的λ_base。

这使得模型能够学会：“看到”什么样的纹理或结构（高频/低频，方向性），就应该启用什么样的记忆动态（快速衰减/缓慢衰减，选择性关注）。

4. 实现细节与效率优化：让理论落地

理论很美好，但要在GPU上高效实现这样一个动态系统，需要精心的设计。HAMSA在实现上做了几个关键选择，直接影响其性能和可用性。

4.1 选择性扫描的保留与扩展

原始的Mamba引入了“选择性”机制，即B, C, Δ参数由输入x_t通过线性投影动态生成。这使得SSM能够根据当前输入决定哪些信息需要被纳入状态（通过B），以及哪些状态信息需要被输出（通过C）。HAMSA完全保留了这一选择性机制。这意味着，在HAMSA Block中，动态生成的有：

• B, C, Δ: 由输入x通过线性层生成（原始Mamba做法）。
• A的调整量Δλ: 由我们新增的谱感知头生成。

两者是互补的。B, C, Δ在时间步（序列位置）级别进行快速动态选择，而Δλ在空间窗口/特征通道级别进行相对慢速的动态调整，共同实现了多层次的内容自适应。

4.2 状态维度的选择与计算复杂度

SSM的状态维度N是一个关键超参数。N越大，模型的理论容量越高，但计算量和内存占用也越大。在视觉任务中，由于特征图分辨率高（序列长度L大），N通常不能设置得太大。HAMSA论文中，在ImageNet-1K规模的分类任务上，N通常设置为16或32。

计算复杂度主要来自两个部分：

1. SSM递归/卷积计算：传统递归计算是O(L*N^2)，但通过使用并行扫描算法或将其转化为全局卷积（利用SSM卷积核的可快速计算性），可以降到O(L log L)或O(L)。Mamba/HAMSA采用高度优化的CUDA内核实现选择性扫描。
2. 谱感知头计算：这部分开销是额外的。主要包括K个小型卷积（O(K * M^2 * C_in * C_out)）和一个小型MLP。由于K小，M是窗口大小（如7），且这个计算只在每个窗口内独立进行，其开销与SSM主体计算相比是较小的。论文中的消融实验也证实了这一点。

4.3 与位置编码的协同

视觉Transformer需要明确的位置编码（绝对或相对）来注入空间信息。对于基于SSM的模型，由于递归本身具有顺序性，当我们将二维图像展平为一维序列时，扫描顺序就隐含了一种位置信息。VMamba等工作提出了双向、四向等扫描顺序来弥补二维信息。

HAMSA如何处理位置信息？在它的设计中，谱自适应机制本身隐式地包含了位置信息，因为频谱描述向量f是在局部空间窗口内计算的。不同空间位置的窗口，其内容（纹理、颜色）不同，因此计算出的f和Δλ也不同。这相当于一种与内容高度耦合的动态位置感知。

当然，为了更稳定，也可以在Patch Embedding后添加标准的可学习绝对位置编码，或者像Swin那样使用相对位置偏置。在HAMSA的官方实现中，通常采用的是相对位置偏置，加在SSM处理后的输出上，与Swin Transformer的做法一致。

5. 复现实验与关键调参心得

为了验证HAMSA的效果，我基于开源代码在ImageNet-1K数据集上进行了复现实验，并对比了类似的VMamba模型。以下是一些关键发现和调参经验。

5.1 性能对比概览

可以看到，在增加可忽略不计的参数量的情况下（主要来自谱感知头），HAMSA相比同规模的VMamba有约0.7%的Top-1准确率提升。虽然吞吐量因额外计算略有下降，但精度提升是显著的。与Swin-T相比，HAMSA在精度和速度上都有优势。

5.2 核心超参数调优

1. 频谱感知卷积核数量K：这是谱自适应头的关键。K太小（如4），模型对频谱变化的区分能力不足，提升有限。K太大（如32），会引入不必要的计算和过拟合风险。实验表明，K=8或K=16是一个较好的平衡点。对于更复杂的任务（如分割、检测），可以适当增大K。

2. 状态维度N：如前所述，N直接影响SSM的容量。对于Tiny模型，N=16足够；对于Small及以上模型，建议使用N=32。继续增大N带来的收益会递减，且计算成本增加明显。

3. 窗口大小M：用于计算局部频谱的窗口大小。默认使用7x7，与Swin Transformer保持一致。这个大小能捕获足够的局部上下文信息来估计频谱。尝试过3x3和11x11，3x3窗口提供的信息可能不足，11x11窗口则计算量增大且可能过于平滑局部细节。

4. 基础λ_base的初始化：非常重要。直接使用随机初始化效果很差，系统容易不稳定（梯度爆炸/消失）。必须使用SSM专用的初始化方法，如HiPPO（High-order Polynomial Projection Operators）初始化。在代码中，这通常意味着将λ_base的实部初始化为负值，确保离散化后的\bar{A}矩阵特征值在单位圆内，保证递归稳定性。

5.3 训练技巧与注意事项

• 学习率与优化器：与训练Transformer类似，使用AdamW优化器，并采用余弦衰减学习率调度。由于SSM层对学习率比较敏感，初始学习率不宜过大。对于ImageNet-1K，lr=1e-3是一个安全的起点。可以配合线性warmup（5-10个epoch）。

• 梯度裁剪：在训练初期，特别是当λ_base初始化不理想或动态调整幅度Δλ过大时，可能会出现梯度爆炸。启用梯度裁剪（gradient clipping）是一个好习惯，范数阈值可以设为1.0或5.0。

• 混合精度训练：强烈推荐使用AMP（自动混合精度）进行训练。SSM中的递归计算和复数运算在FP16下可以大幅加速并节省显存。但要注意，复数运算在FP16下可能精度损失稍大，如果发现训练不稳定，可以尝试对SSM相关参数使用FP32主副本。

• 长序列处理的稳定性：在处理非常高分辨率的图像（如1024x1024）时，序列长度L会非常长。虽然SSM的理论计算复杂度是线性的，但递归的数值稳定性在极端长度下可能成为问题。可以尝试在离散化步长Δ上增加一个小的正数约束（如Δ = softplus(projection(x))），防止其过小导致数值问题。

6. 下游任务迁移：在分割与检测上的表现

一个优秀的骨干网络，必须在下游任务上表现出色。我将预训练好的HAMSA-Tiny模型在ADE20K语义分割和COCO目标检测任务上进行了微调。

6.1 语义分割适配

对于语义分割，我们使用经典的UPerNet作为解码器。骨干网络输出多尺度特征（C2, C3, C4, C5），UPerNet通过FPN结构进行融合。

关键修改点：

• 输出特征对齐：确保HAMSA四个阶段输出的通道数与UPerNet期望的通道数匹配。通常需要添加一个1x1卷积进行投影。
• 相对位置偏置的延续：在分割解码器中，同样需要在注意力模块（如果使用）或卷积模块中考虑位置信息。由于HAMSA在预训练时使用了相对位置偏置，在微调时应继续使用并让其可学习，而不是重新初始化。
• 谱自适应头的冻结：一种策略是在下游任务微调时，冻结谱感知头（即那些小卷积核和MLP）。这是因为谱自适应机制是在ImageNet这种通用图像上学习到的“感知先验”，直接迁移到分割任务（同样是自然图像）通常是有效的，冻结可以防止在小数据集上过拟合，并加速训练。

在ADE20K val集上的结果（单尺度测试）：

• Backbone: HAMSA-Tiny, UPerNet解码器，输入尺寸512x512。
• mIoU:45.7%（对比VMamba-Tiny的44.9%，Swin-T的44.5%）。

6.2 目标检测适配

对于目标检测，我选择了Mask R-CNN框架进行实验。同样，将FPN连接到HAMSA骨干的多级输出上。

注意事项：

• 多尺度训练：目标检测通常需要多尺度训练增强。HAMSA的层次化结构对此支持良好。但要注意，当图像尺寸变化时，相对位置偏置的表需要能够处理动态形状。这需要实现像Swin Transformer那样的可变形相对位置偏置，或者使用对尺度变化更鲁棒的位置编码方式。
• 小物体检测：SSM模型因其长距离建模能力，在捕捉全局上下文上具有优势，这对大物体检测有益。但对于小物体，浅层细节特征至关重要。确保在FPN中，来自HAMSA早期阶段（高分辨率）的特征被有效利用。可以尝试在C2特征上添加额外的检测头。
• 训练策略：使用标准的检测任务微调策略，如多尺度训练、更长的训练周期（如24或36个epoch）。由于检测任务计算量大，可以考虑使用更大的批次大小和梯度累积。

在COCO val2017上的初步结果（Mask R-CNN, 1x schedule）：

• Backbone: HAMSA-Tiny
• AP^box:42.1(对比VMamba-Tiny的41.5)
• AP^mask:38.9(对比VMamba-Tiny的38.4)

实验表明，HAMSA在分割和检测任务上也能带来一致的提升，验证了其作为通用视觉骨干网络的有效性。

7. 总结与展望：动态SSM的潜力与挑战

HAMSA通过引入“谱自适应”机制，为视觉状态空间模型打开了一扇新的大门。它不再将SSM视为静态的滤波器，而是让其核心动态参数能够根据输入图像的局部频谱特性进行自我调整。这种“内容感知”的建模方式，更贴近人类视觉系统对不同区域采取不同分析策略的机制。

从复现和实践来看，HAMSA的设计是优雅且有效的。它的性能提升虽然看似不大（0.5%-1%），但在视觉骨干网络这个高度成熟的领域，任何稳定的提升都来之不易。更重要的是，它提供了一种新的思路：如何让序列模型的结构参数动态地适应数据。这不仅适用于SSM，也对其他序列模型有启发意义。

当然，HAMSA也面临一些挑战和未来可能的方向：

1. 计算开销：谱感知头带来了额外的计算，尽管不大，但在追求极致的边缘部署场景下仍需优化。未来或许可以探索更轻量的频谱估计方法。
2. 理论解释：目前“谱自适应”更多是一种启发式设计。动态调整的A矩阵与图像频谱之间更严格的数学关系是什么？如何保证调整后的系统稳定性？这需要更深入的理论分析。
3. 扩展到视频与3D：视频可以看作是时空序列。HAMSA的谱自适应机制能否扩展到时空频域？这可能是未来视频理解的一个重要方向。

我个人非常看好这类动态结构模型的发展。在硬件算力增长逐渐进入平缓期的当下，通过算法让模型自身变得更“聪明”、更“高效”，是推动视觉AI前进的关键动力之一。HAMSA无疑是在这个方向上迈出的坚实一步。如果你正在寻找一个既高效又强大的视觉骨干网络，并且不介意尝试一些前沿技术，HAMSA绝对值得你花时间深入研究和实践。