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架构全景层次化设计如何融入动态SSMHAMSA的整体架构遵循了现代视觉骨干网络的设计范式即层次化金字塔结构。这种结构的好处显而易见浅层捕获细节和边缘深层捕获语义和物体。对于下游任务如目标检测、语义分割多尺度特征图是必不可少的。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来微调每个通道的λ_iHAMSA通过一个小的MLP多层感知机来实现Δλ MLP_Spectrum(f)这个MLP以f为输入输出一个N维的复数向量Δλ ∈ C^N。然后我们将这个调整量加到基础或共享的A矩阵参数上λ_adapted λ_base Δλ这里λ_base是所有HAMSA Block共享的、可学习的基础参数初始值可以采用HiPPO等经典方法。而Δλ则是每个窗口、每个样本、每个网络层都不同的动态调整部分。3.3 自适应过程的计算图与梯度流理解这个过程的梯度流很重要。λ_base是全局可学习参数其梯度来自于所有样本和所有位置的反向传播。而Δλ是通过一个由频谱描述向量f参数化的网络生成的f又依赖于当前输入特征X和那K个可学习的频谱感知卷积核。因此整个系统形成了一个闭环输入特征X决定了局部频谱特性f。f通过MLP决定了SSM动态系统参数A的调整量Δλ。调整后的SSM (A_adapted) 处理特征X产生输出。输出的损失梯度会同时更新特征提取权重、频谱感知卷积核、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。计算复杂度主要来自两个部分SSM递归/卷积计算传统递归计算是O(L*N^2)但通过使用并行扫描算法或将其转化为全局卷积利用SSM卷积核的可快速计算性可以降到O(L log L)或O(L)。Mamba/HAMSA采用高度优化的CUDA内核实现选择性扫描。谱感知头计算这部分开销是额外的。主要包括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 性能对比概览模型参数量ImageNet Top-1 Acc吞吐量 (imgs/s)VMamba-Tiny28M82.5%1250HAMSA-Tiny29M83.2%1180Swin-T28M81.3%755VMamba-Small50M83.5%980HAMSA-Small52M84.1%920可以看到在增加可忽略不计的参数量的情况下主要来自谱感知头HAMSA相比同规模的VMamba有约0.7%的Top-1准确率提升。虽然吞吐量因额外计算略有下降但精度提升是显著的。与Swin-T相比HAMSA在精度和速度上都有优势。5.2 核心超参数调优频谱感知卷积核数量K这是谱自适应头的关键。K太小如4模型对频谱变化的区分能力不足提升有限。K太大如32会引入不必要的计算和过拟合风险。实验表明K8或K16是一个较好的平衡点。对于更复杂的任务如分割、检测可以适当增大K。状态维度N如前所述N直接影响SSM的容量。对于Tiny模型N16足够对于Small及以上模型建议使用N32。继续增大N带来的收益会递减且计算成本增加明显。窗口大小M用于计算局部频谱的窗口大小。默认使用7x7与Swin Transformer保持一致。这个大小能捕获足够的局部上下文信息来估计频谱。尝试过3x3和11x113x3窗口提供的信息可能不足11x11窗口则计算量增大且可能过于平滑局部细节。基础λ_base的初始化非常重要。直接使用随机初始化效果很差系统容易不稳定梯度爆炸/消失。必须使用SSM专用的初始化方法如HiPPOHigh-order Polynomial Projection Operators初始化。在代码中这通常意味着将λ_base的实部初始化为负值确保离散化后的\bar{A}矩阵特征值在单位圆内保证递归稳定性。5.3 训练技巧与注意事项学习率与优化器与训练Transformer类似使用AdamW优化器并采用余弦衰减学习率调度。由于SSM层对学习率比较敏感初始学习率不宜过大。对于ImageNet-1Klr1e-3是一个安全的起点。可以配合线性warmup5-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, C5UPerNet通过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 scheduleBackbone: HAMSA-TinyAP^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也面临一些挑战和未来可能的方向计算开销谱感知头带来了额外的计算尽管不大但在追求极致的边缘部署场景下仍需优化。未来或许可以探索更轻量的频谱估计方法。理论解释目前“谱自适应”更多是一种启发式设计。动态调整的A矩阵与图像频谱之间更严格的数学关系是什么如何保证调整后的系统稳定性这需要更深入的理论分析。扩展到视频与3D视频可以看作是时空序列。HAMSA的谱自适应机制能否扩展到时空频域这可能是未来视频理解的一个重要方向。我个人非常看好这类动态结构模型的发展。在硬件算力增长逐渐进入平缓期的当下通过算法让模型自身变得更“聪明”、更“高效”是推动视觉AI前进的关键动力之一。HAMSA无疑是在这个方向上迈出的坚实一步。如果你正在寻找一个既高效又强大的视觉骨干网络并且不介意尝试一些前沿技术HAMSA绝对值得你花时间深入研究和实践。