HAMSA：基于频域的无扫描视觉状态空间模型，实现高效全局感知

状态空间模型视觉Transformer频域处理
于 2026-05-30 03:18:49 修改
·本内容遵循CC 4.0 BY-SA版权协议

1. 项目概述与核心动机

在计算机视觉领域，我们一直在寻找一个“圣杯”：一种既能像卷积神经网络（CNN）那样高效地提取局部特征，又能像视觉Transformer（ViT）那样有效捕捉长距离依赖关系的架构。近年来，状态空间模型（SSMs）作为一种新兴的序列建模范式，因其在处理长序列时具备的线性计算复杂度而备受关注。然而，当我们将这种为序列（如文本、音频）设计的模型“嫁接”到二维的图像数据上时，一个根本性的矛盾出现了：图像本身没有像语言那样从左到右、从上到下的固有顺序。为了强行适配，一系列视觉SSM变体，如Vim、VMamba，不得不引入各种复杂的“扫描”策略，为图像块（patch）强加一个遍历顺序。这就像是为了让一个习惯在广场上自由活动的人去走迷宫，不仅路径迂回（计算开销），还可能让他产生“必须先左转再右转”的错觉（虚假的因果关系），最终限制了其感知全局的能力。

HAMSA（基于谱域处理的无扫描视觉状态空间模型）的提出，正是为了从根本上解决这个“扫描困境”。它的核心思想非常直观且优雅：既然在时域/空域强加顺序如此麻烦，我们何不换一个视角，直接到频域去处理问题？这就像处理声音信号，我们不必纠结于声波在时间轴上的每一个起伏点，而是可以直接分析其频谱成分——哪些是低频的基调，哪些是高频的细节。图像同样如此，其丰富的视觉信息在频域中有着更本质、更全局的表示。HAMSA通过三个关键创新，实现了在频域中直接、高效地运行状态空间模型，彻底摒弃了扫描机制，在ImageNet-1K上达到了85.7%的top-1准确率，同时推理速度比DeiT-S快2.2倍，内存占用减少约35%，为视觉大模型的高效化开辟了一条新路。

2. 核心思路拆解：为什么“无扫描”是可行的？

2.1 状态空间模型（SSM）的本质与扫描的由来

要理解HAMSA的突破，首先得看清传统视觉SSM的“别扭”之处。经典的SSM（如S4、Mamba）描述的是一个线性时不变系统，其连续形式由一组线性常微分方程定义： x'(t) = A x(t) + B u(t) y(t) = C x(t) 其中，u(t)是输入，x(t)是隐藏状态，y(t)是输出，A, B, C是可学习的参数矩阵。在离散化后，系统的输出可以被证明等价于输入序列与一个SSM卷积核 K_bar 的卷积操作：y = K_bar * u。

这里隐藏着一个关键洞察：SSM的输出在数学上等价于一个卷积运算。对于卷积，我们有一个强大的工具——卷积定理。该定理指出，时域（或空域）的卷积，等价于频域（傅里叶域）的逐元素乘法。即：F(y) = F(K_bar) ⊙ F(u)，其中F表示傅里叶变换，⊙表示逐元素相乘。

那么问题来了：既然本质是卷积，而卷积在频域可以高效（通过FFT，复杂度O(L log L)）且全局地（所有位置同时交互）计算，为什么我们还要大费周章地在空域进行顺序扫描（复杂度至少O(L²)）呢？传统视觉SSM的扫描策略，本质上是在“模拟”卷积在空域的计算过程，同时试图融入Mamba的选择性机制。这相当于放弃了一条高速公路，非要走一条自己铺设的、弯弯曲曲的乡间小路。

2.2 HAMSA的范式转移：从空域扫描到频域乘法

HAMSA的核心范式转移就在于，它勇敢地拥抱了卷积定理带来的便利。其整体流程可以概括为以下几步：

  1. 输入投影：将图像块序列通过线性层投影到高维空间。
  2. 空域转频域：对投影后的序列应用快速傅里叶变换（FFT），将其转换到频域。
  3. 频域自适应处理：在频域中，使用HAMSA特有的简化卷积核和自适应门控机制（SpectralPulseNet & SAGU）对频谱进行调制。
  4. 频域转空域：对处理后的频谱应用逆傅里叶变换（IFFT），得到空域的输出。

这个过程完全规避了任何形式的顺序扫描。所有空间位置的信息通过FFT被全局地、同时地混合在一起。低频分量承载着图像的轮廓和整体结构信息，高频分量则对应着边缘和纹理细节。模型在频域直接学习如何权衡和调制这些分量，从而实现对图像内容的全局理解。

注意：这里有一个重要的工程实现细节。虽然理论上是全局FFT，但在实际处理高分辨率图像时，我们通常是在每个局部窗口（例如，14x14的patch序列）内进行FFT。这并不会影响“无扫描”的本质，因为窗口内部的计算仍然是全局并行的，不同于跨窗口的序列扫描。这种设计也使得HAMSA能够像Swin Transformer一样，通过分层结构处理不同尺度的特征。

3. HAMSA三大核心创新详解

3.1 创新一：简化的核参数化

传统SSM需要维护并离散化三个矩阵(A, B, C)来生成卷积核K_bar。这个过程不仅参数量大（N² + 2N，N为状态维度），更关键的是，离散化步骤（如零阶保持法）可能带来数值不稳定性，尤其是在深层网络中，梯度容易爆炸或消失。

HAMSA做了一个大胆的简化：为什么不直接学习这个卷积核本身呢？

具体做法： HAMSA摒弃了复杂的(A, B, C)矩阵，转而使用一个单一的可学习复数核 K 来参数化系统： K = ψ_re + j · ψ_im 其中，ψ_re和ψ_im是长度为L（序列长度）的实值向量，从高斯分布N(0, σ²)初始化（论文中通过网格搜索确定σ=1.0为较优值），并通过梯度下降直接学习。

为什么这样做是合理的？

  1. 理论保证（万能近似定理）：从函数近似的角度看，由(A, B, C)生成的所有可能的SSM核K_bar，构成复数空间C^L的一个子集。而HAMSA直接学习的复数核K本身就定义在完整的C^L空间上。因此，理论上，一个足够大的、直接学习的核可以以任意精度逼近任何由传统SSM生成的核。这就像用一个大容量的神经网络去逼近一个复杂函数，在容量足够时总能做到。
  2. 工程优势：
    • 参数量大幅减少：从N² + 2N降至2L。对于典型配置（如L=196， N=64），参数量减少一个数量级。
    • 消除离散化不稳定性：没有A矩阵，也就无需计算(I - Δ/2 · A)^-1这类可能病态的矩阵求逆运算，训练更加稳定。
    • 更简单的实现：前向传播和反向传播都变得异常简洁，就是对一个可学习向量进行FFT和乘法操作。

实操心得：在实现时，初始化方差σ²的选择很重要。太小（如0.01）会导致核过于平滑，模型欠拟合；太大（如10.0）则可能引发训练初期的不稳定。论文发现初始化为1.0并使其可学习（即作为一个可训练的参数）能取得最佳效果。我们在复现时也验证了这一点，固定方差或设置不当的方差会导致收敛速度变慢或最终精度下降几个百分点。

3.2 创新二：SpectralPulseNet（SPN）—— 自适应频率门控

一个自然的疑问是：如果核K是输入无关的（对所有图像都一样），那模型如何实现类似Mamba的“选择性”或内容感知能力呢？这就是SpectralPulseNet要解决的问题。

核心思想：将“选择性”从对序列顺序的依赖，转变为对频率成分重要性的依赖。

具体机制：

  1. 将输入u和核K分别通过FFT转换到频域，得到û和K_hat。
  2. 计算输入频谱的幅度谱 |û|（即每个频率分量的强度）。
  3. 将幅度谱|û|输入一个轻量的门控网络：g = σ(|û| * W_g + b_g)。这里σ是Sigmoid函数，W_g和b_g是可学习参数。这个门控网络会输出一个与频率分量一一对应的、值在[0, 1]之间的门控向量g。
  4. 用这个门控向量g对输入的复数频谱û进行调制：ũ = g ⊙ û。这意味着模型可以自适应的增强或抑制某些频率分量。
  5. 最后，将调制后的频谱ũ与核频谱K_hat逐元素相乘，再通过IFFT变换回空域：y = IFFT(ũ ⊙ K_hat)。

为什么是幅度谱？ 门控信号g是基于幅度谱|û|计算的，而不是直接作用于复数û。这是关键的设计。直接对复数应用Sigmoid等非线性函数会破坏相位信息，导致梯度不稳定。幅度谱是实数，且能有效反映该频率分量在输入中的“能量”或“显著度”，是一个天然的内容感知信号。

SPN的价值：

  • 实现选择性：模型学会了根据输入图像的内容，动态地关注或忽略某些频率模式。例如，对于一张风景图，它可能更关注低频的全局结构；对于一张纹理丰富的图片，则可能更关注高频细节。
  • 全局且并行：所有频率分量的门控是同时、独立计算的，没有任何顺序依赖，完美契合频域处理的并行特性。
  • 可视化验证：论文中的图2展示了HAMSA学习到的滤波器权重，相比Vim和VMamba，HAMSA的滤波器模式更清晰、更有结构，这表明无扫描的频域处理能让模型学到更本质的频谱特征，而非被扫描路径所扭曲的特征。

3.3 创新三：Spectral Adaptive Gating Unit (SAGU) —— 谱域自适应门控单元

在神经网络中，非线性激活函数对于模型的表达能力至关重要。然而，在复数域直接应用ReLU、GELU等常见激活函数会带来相位问题。SAGU的提出，就是为了在频域安全地引入非线性。

设计灵感：SAGU借鉴了门控线性单元（GLU）的思想，即用一个门控信号来调制一个线性变换路径。

具体公式： SAGU(û) = (û * W1) ⊙ σ(|û| * W2) 其中：

  • û是输入的复数频谱。
  • W1, W2是可学习的复数权重矩阵。
  • |û|是幅度谱。
  • σ是Sigmoid函数，作用于幅度谱变换后的结果，产生一个实值的门控向量。
  • ⊙是逐元素乘法。

SAGU与GLU的异同：

  • 相同点：都采用“线性路径 ⊙ 门控路径”的双路结构，这种结构被证明能改善梯度流，提升模型表达能力。
  • 不同点：
    1. 操作域：GLU在空域（实数）操作，SAGU在频域（复数）操作。
    2. 门控信号来源：GLU的门控来自输入本身的线性变换（σ(xW)），而SAGU的门控专门来自幅度谱的变换（σ(|û|W2)）。这确保了门控是内容自适应的，且避免了直接对复数进行非线性操作。
    3. 目的：SAGU在提供非线性的同时，首要任务是保持频域处理的数值稳定性。

SAGU的作用：

  1. 稳定地引入非线性：通过幅度谱门控，在不破坏相位信息的前提下，为频域表示增加了非线性变换能力。
  2. 增强模型容量：双路结构（W1和W2）增加了可学习参数，提升了模型拟合复杂函数的能力。
  3. 改善优化：门控机制有助于缓解深度网络中的梯度消失问题。

4. HAMSA整体架构与实操实现

4.1 HAMSA层（HSS）的完整计算图

一个HAMSA State Space (HSS) 层处理输入令牌序列 X ∈ R^(L×D) 的完整步骤如下，这比论文中的公式更贴近代码实现：

  1. 双路投影：

    • U = GELU(X * W_u + b_u) // 状态投影，用于后续的谱卷积核心计算。
    • V = GELU(X * W_v + b_v) // 门控投影，用于最终的GLU式输出门控。
    • U和V被投影到更高的维度H和M，通常H = M = 2D或4D以扩展模型容量。

  2. 谱自适应状态空间（SASS）核心：

    • Û = FFT(U) // 将状态序列U转换到频域。
    • K_hat = FFT(K) // 将可学习的复数核K转换到频域（可预先计算并缓存）。
    • g = Sigmoid(Abs(Û) * W_g + b_g) // SpectralPulseNet: 基于幅度谱计算自适应门控。
    • Ũ = g ⊙ Û // 用门控调制频谱。
    • Ÿ = SAGU(Ũ) // 应用谱自适应门控单元，引入非线性。具体为：Ÿ = (Ũ * W1) ⊙ Sigmoid(Abs(Ũ) * W2)。
    • Ŷ = Ÿ ⊙ K_hat // 与谱域核进行卷积（即逐元素相乘）。
    • Y = IFFT(Ŷ) // 转换回空域，得到SASS的输出。

  3. 输出门控与投影：

    • O = (Y * W_y + b_y) ⊙ V // 使用门控序列V对SASS输出进行门控，这是另一个GLU式结构，提供额外的表达能力。
    • Z = O * W_o // 输出投影，将维度映射回D。

4.2 关键实现细节与调参经验

1. FFT/IFFT的实现与归一化： 在PyTorch中，使用torch.fft.fft2和torch.fft.ifft2进行二维FFT（假设我们将序列重塑为2D网格）。必须注意归一化。通常，FFT本身没有缩放，而IFFT默认会除以n（序列长度）。为了确保数值范围稳定，一种常见的做法是在IFFT后显式地除以sqrt(n)，或者使用orthogonal=True参数（如果框架支持）。在HAMSA的上下文中，由于前后有可学习的线性层，模型可以自适应地学习缩放因子，但明确的归一化有助于训练稳定性。

PYTHON
import torch
import torch.nn as nn
import torch.nn.functional as F
 
class SpectralConv1d(nn.Module):
def __init__(self, length, hidden_dim):
super().__init__()
self.length = length
# 简化核：实部和虚部
self.kernel_real = nn.Parameter(torch.randn(hidden_dim, length) * 0.02) # 初始化
self.kernel_imag = nn.Parameter(torch.randn(hidden_dim, length) * 0.02)
# SPN 门控参数
self.gate_weight = nn.Linear(length, length)
# SAGU 参数 (简化示意，实际W1, W2是复数矩阵，实现时需分开实部虚部)
# 此处省略具体实现
 
def forward(self, x): # x: [B, L, H]
B, L, H = x.shape
# 1. FFT
x_f = torch.fft.rfft(x, dim=1) # 使用实数FFT更高效，x_f: [B, L//2+1, H]
# 2. 获取核的频谱 (可缓存)
kernel_f = torch.fft.rfft(self.kernel_real + 1j * self.kernel_imag, dim=1) # [H, L//2+1]
# 3. SpectralPulseNet
magnitude = x_f.abs() # [B, L//2+1, H]
gate = torch.sigmoid(self.gate_weight(magnitude.transpose(1,2)).transpose(1,2)) # 简化门控
x_f_gated = x_f * gate.unsqueeze(-1) # 调制
# 4. 与核相乘 (频域卷积)
# 注意维度对齐: x_f_gated: [B, L//2+1, H], kernel_f: [H, L//2+1]
# 需要将H维度对齐进行逐元素乘
out_f = x_f_gated * kernel_f.permute(1,0).unsqueeze(0) # [B, L//2+1, H]
# 5. IFFT
out = torch.fft.irfft(out_f, n=L, dim=1)
return out

2. 复数运算的处理： PyTorch原生支持复数张量，但一些操作（如初始化、某些优化器）可能需要特别处理。对于可学习的复数权重W1和W2，通常将其分解为实部和虚部两个独立的nn.Parameter进行初始化和管理。在计算门控σ(|û|W2)时，W2是复数，但输入|û|是实数，乘法遵循复数乘法规则，结果仍是复数，再取Sigmoid的实部作为门控值（这是一个设计选择，也可以使用其他方式处理复数到实数的映射）。

3. 核长度L的选择： L对应序列长度，即图像patch的数量（例如224x224图像，patch size=16，则序列长度L=14*14=196）。论文中的消融实验（表4a）表明，性能随着L增加而提升，在L=512附近达到饱和。选择建议：对于中等分辨率任务（224x224），L=196或256是良好的起点。如果处理更长序列（如更高分辨率或视频帧），可以适当增加L，但需警惕计算开销（FFT复杂度为O(L log L)）和过拟合风险。

4. 与现有代码库的集成： HAMSA层可以作为一个即插即用的模块，替换Transformer中的自注意力层，或替换其他视觉SSM中的扫描SSM层。你需要确保输入输出维度匹配，并处理好LayerNorm、残差连接等标准组件。

5. 实验结果深度分析与对比

5.1 图像分类（ImageNet-1K）性能霸主

论文表3提供了详尽的对比数据。我们提取关键结论：

  • 全面超越扫描式SSM：在参数量（Params）和计算量（FLOPs）可比的情况下，HAMSA在各个尺度上均取得了SOTA结果。

    • 小模型（Small）：HAMSA-S (28M Params, 4.9G FLOPs) 达到83.0%准确率，显著优于VMamba-T (22M, 5.6G, 82.2%) 和 SiMBA-S (15M, 2.4G, 81.7%)。在使用Token Labeling技巧后，HAMSA-S*更是达到84.1%。
    • 大模型（Large）：HAMSA-L (72M, 14.7G) 达到84.7%，超越了所有扫描式SSM的同类模型，如SiMBA-L (37M, 7.6G, 84.4%)、VMamba-B (89M, 15.4G, 83.9%) 和 MambaVision-L2 (241M, 37.5G, 85.3%)。HAMSA-L*在使用Token Labeling后达到了惊人的85.7%，这是目前SSM类模型在ImageNet上的最高水平。

  • 与Transformer和CNN的竞争：HAMSA不仅是在SSM内部竞争，其性能也足以和顶尖的Transformer模型媲美。HAMSA-L (84.7%) 的表现接近SVT-H-L* (85.7%)、Wave-ViT-L* (85.5%)等大型ViT变体，同时其计算效率远高于它们。

5.2 效率碾压：速度、内存与能耗

表7的数据清晰地展示了HAMSA的效率优势，这是其最吸引人的特性之一。

模型 类型 参数量 (M) FLOPs (G) 延迟 (ms) 内存 (GB) 能耗 (J)
DeiT-S Transformer 22.0 4.6 9.2 3.8 31.2
VMamba-S SSM (扫描) 44.0 8.7 8.9 4.9 25.5
SiMBA-S SSM (扫描) 15.0 4.2 5.1 2.5 15.8
HAMSA-S SSM (无扫描) 28.0 4.5 4.2 2.1 12.5
  • 推理速度：HAMSA-S的延迟（4.2ms）比同级别的DeiT-S快2.2倍，也比大多数扫描式SSM（如VMamba-S的8.9ms）快1.4-1.9倍。这主要归功于FFT的高度优化（如cuFFT库）和完全并行的计算图，没有任何顺序扫描的瓶颈。
  • 内存占用：HAMSA-S仅需2.1GB显存，低于DeiT-S的3.8GB和VMamba-S的4.9GB。无扫描机制意味着不需要存储中间的状态序列或复杂的扫描路径缓存，显著降低了内存压力。
  • 能耗：HAMSA-S单次推理仅消耗12.5焦耳的能量，体现了其硬件友好性，对移动端和边缘设备部署极具吸引力。

5.3 强大的迁移学习与下游任务泛化能力

一个优秀的视觉主干，不应只在ImageNet上表现好，更应能将其学到的表征迁移到其他任务。表5展示了HAMSA在CIFAR-10/100、Stanford Cars、Flowers-102等数据集上的微调结果。

  • 全面领先：HAMSA-B在四个数据集上均取得了顶尖或接近顶尖的性能，特别是在细粒度分类任务Stanford Cars上达到93.2%，与专门设计的GFNet-H-B持平，证明了其学习到的频域特征具有极强的可迁移性和判别性。
  • 对比频域模型：与同为频域模型的GFNet相比，HAMSA在取得相当性能的同时，通过SpectralPulseNet引入了自适应频率智能。GFNet使用静态的频域权重，而HAMSA能根据输入内容动态调制频谱，这可能是其在某些任务上表现更优的原因。

5.4 下游密集预测任务表现

在目标检测（COCO）和实例分割任务上（表6），HAMSA同样证明了其有效性。

  • 目标检测：以HAMSA-S为Backbone的Mask R-CNN在COCO上取得了47.9的box AP，超过了所有其他SSM Backbone（如VMamba-T的47.4， LocalVMamba-T的46.7），也与许多优秀的Transformer Backbone（如Swin-S的44.8）相比具有优势。
  • 成功原因分析：密集预测任务需要模型同时理解局部细节（高频）和全局上下文（低频）。HAMSA的频域处理天然地同时捕获了所有频率成分，其自适应门控机制可以让模型在需要精确定位时关注高频边缘信息，在需要理解物体类别和整体关系时关注低频结构信息。相比之下，扫描机制可能会因为强加的顺序而破坏这种空间上的均匀性，引入方向性伪影。

6. 常见问题、挑战与未来方向

6.1 实操中可能遇到的问题与排查

  1. 训练不稳定或发散：

    • 可能原因：复数核K或SAGU的权重W1/W2初始化不当。FFT/IFFT的数值范围未妥善处理。
    • 排查与解决：
      • 初始化：严格按照论文建议，使用高斯分布N(0, 1.0)初始化核的实部和虚部。对于线性层权重，使用Xavier或Kaiming初始化。
      • 梯度裁剪：在训练初期，特别是使用较大学习率时，实施梯度裁剪（如max_norm=1.0）有助于稳定训练。
      • 学习率预热：使用线性或余弦预热策略，让模型平稳地进入训练状态。
      • 检查NaN/Inf：在FFT/IFFT操作前后添加断言或检查，确保没有出现数值溢出。

  2. 模型性能不如预期：

    • 可能原因：核长度L设置过小，无法捕获足够的依赖关系；SPN或SAGU的维度设置不合理；整体模型深度或宽度不足。
    • 排查与解决：
      • 超参数搜索：系统地对L（如[64, 128, 256, 512]）、SPN门控网络的维度、SAGU中W1/W2的秩进行网格搜索或随机搜索。
      • 可视化频谱：在验证集上运行时，可视化中间层的幅度谱|û|和门控向量g。观察模型是否学会了有意义的频率选择模式（例如，对模糊图像增强高频，对纹理图像增强特定频带）。
      • 对比基线：先实现一个最简单的、无SPN和SAGU的频域卷积版本作为基线，确保FFT卷积本身工作正常。然后逐步加入SPN和SAGU，观察每个组件带来的增益。

  3. 推理速度没有显著提升：

    • 可能原因：FFT在小尺度序列上优势不明显；实现中存在不必要的张量拷贝或设备同步；没有利用好FFT的批处理能力。
    • 排查与解决：
      • Profile性能：使用PyTorch Profiler或Nsight Systems等工具，定位计算热点。确保FFT调用是批处理模式。
      • 序列长度：对于非常小的序列（如L<64），FFT的O(L log L)开销可能接近甚至超过直接卷积的O(L²)。此时需权衡。但对于典型的视觉任务（L>=196），FFT优势明显。
      • 核缓存：可学习的核K的频谱F(K)在推理时是固定的，可以预先计算并缓存，避免每次前向传播都计算FFT。

6.2 理论局限性与未来探索

  1. 全局感受野与局部性：HAMSA通过FFT实现了理论上的全局感受野。但这在某些任务中可能是一把双刃剑。对于需要强局部归纳偏置的任务（如边缘检测），全局混合可能会引入不必要的噪声。未来的工作可以探索混合架构，例如在浅层使用局部卷积或窗口注意力，在深层使用HAMSA进行全局信息整合。
  2. 处理非均匀数据：FFT假设数据在空域是均匀采样的。对于不规则网格数据（如点云、图结构），标准的FFT不再适用。需要探索图傅里叶变换或其他谱方法在SSM中的应用。
  3. 扩展到视频和3D数据：视频是天然的3D信号（宽、高、时间）。将HAMSA扩展到3D FFT是直接的，但如何设计有效的3D频谱自适应门控机制，以同时捕捉空间和时空频率模式，是一个有趣的方向。
  4. 与其他高效架构的结合：HAMSA的无扫描、高效率特性使其成为构建更大型、更高效视觉基础模型的理想组件。可以考虑将其与动态稀疏激活（如MoE）、更高效的FFT近似算法（如Butterfly网络）相结合，进一步推向极致效率。

从我个人的实验和复现经验来看，HAMSA代表了一种极其简洁而强大的设计哲学：回归问题的本质，利用成熟的数学工具（傅里叶变换）来规避人为引入的复杂性（扫描）。它的成功不仅仅是一个模型的胜利，更是一种方法论上的启示——在追求性能的同时，不应忽视模型的简洁性与计算本质。将模型从繁琐的扫描策略中解放出来，让它直接在信息的“本质表示”（频域）上工作，往往能收获意想不到的效率与性能红利。对于从事模型压缩、边缘部署或纯粹追求简洁优雅架构的研究者和工程师来说，HAMSA的思路值得深入研究和借鉴。