1. 为什么Vision Mamba值得单独拎出来聊

视觉骨干网络这个赛道，卷积神经网络统治了将近十年，Transformer进来之后又把注意力机制推上了神坛，但这两条路线各自都带着一个甩不掉的包袱。卷积的感受野受限于堆叠深度，想看得更远就得把网络做得更深更重；Transformer的自注意力是全局的，代价是计算量随序列长度呈平方级增长，处理高分辨率图像时显存和延迟都很难看。Vision Mamba（后面统一简称ViM）就是在这个背景下被提出来的，它把状态空间模型（State Space Model，SSM）那套线性复杂度的序列建模思路搬到了视觉任务上，试图在“看得远”和“算得起”之间找一个平衡点。

我第一次接触ViM的时候，最直观的感受是：这东西的论文公式看着唬人，但核心思想其实很朴素——用一条可以沿空间方向扫描的隐状态，把图像块序列的信息逐步压缩传递下去，从而在不做全局两两注意力计算的前提下，依然让每个位置都能间接“看到”前面所有位置的信息。它解决的问题很明确：在保持接近线性复杂度的同时，获得长距离依赖建模能力。适合谁来读这篇内容？如果你已经对Transformer的注意力机制有基本认知，想搞清楚SSM到底怎么在视觉里落地，或者你正在做高分辨率图像分类、密集预测类任务，被显存和推理速度卡住了脖子，那这篇拆解会对你有直接帮助。哪怕你只是听说过Mamba这个词但一直没搞明白它和Transformer的本质区别，我也会尽量用生活化的类比把这条链路讲透。

需要提前说明的是，ViM这类模型涉及的选择性扫描、硬件感知实现等细节，原始论文给出的是一种方案，实际工程落地时往往还要结合具体框架和硬件做调整。下面我讲到的参数配置、实现步骤，一部分来自论文本身的设定，一部分是基于我在类似序列建模任务中的常见实践做的合理补全，我会在涉及推断的地方明确标注出来，方便你区分哪些是“原味”、哪些是“工程经验”。

2. 核心思路拆解：ViM到底在解决什么问题

2.1 从“两两对比”到“逐级传递”的范式转换

要理解ViM，先得把Transformer的注意力机制想清楚。自注意力的本质是让序列里每个元素和其他所有元素做一次相似度计算，然后按相似度加权聚合信息。这就像开一场所有人都要互相握手的会议，n个人就是n的平方次握手，人一多会场就乱套。图像被切成patch之后，序列长度动辄几百上千，平方级的开销立刻变成瓶颈。

ViM换了个思路：不搞全员握手，而是排成一队，信息从队首往队尾逐个传递，每个人只需要接收前一个人手里的“会议纪要”并更新它。这个“会议纪要”就是隐状态。这样一来，每个人处理的信息量是固定的，总开销随人数线性增长。代价是信息传递有方向性，队尾的人天然比队首的人掌握更多历史信息，所以怎么排队、按什么顺序扫描，就成了设计的关键。这个类比能帮你抓住SSM在视觉里的核心：它是一个带记忆的递归式信息压缩器，而不是一个全局比对器。

2.2 选择性机制：让“记忆”学会该记什么、该忘什么

早期的SSM（比如S4）参数是输入无关的，也就是说不管当前看到的是天空还是车轮，记忆更新的规则都一样，这在语言任务里勉强够用，但在视觉里就太死板了。ViM采用的Mamba结构引入了**选择性（selective）**机制：更新隐状态的几个关键参数（可以理解为“写入门控”“遗忘门控”“输出门控”）都由当前输入动态生成。翻译成人话就是，模型看到平坦的天空区域时倾向于少记、快速略过，看到纹理复杂的物体边缘时倾向于多记、精细保留。

这个设计带来的直接好处是模型具备了内容感知的压缩能力，但副作用也很明显：参数随输入变化之后，原本可以用卷积形式并行计算的固定核就失效了，必须改成递归的串行扫描。这就是为什么Mamba类模型对硬件实现特别敏感，也是后面要讲的“硬件感知扫描”存在的根本原因。理解这一点，你就明白ViM不是一个纯算法层面的优雅公式，它和底层计算模式是深度绑定的。

2.3 视觉场景下的方向性难题

语言是天然的一维序列，从左到右扫一遍就完事。图像是二维的，你把patch拉平成一维，无论怎么排，都会人为制造出“相邻但不相连”的断裂。ViM的处理方式是做双向扫描：正向扫一遍，反向再扫一遍，然后把两个方向的输出融合。这有点像读书时先顺着读一遍再倒着读一遍，两遍下来对全局结构的把握更完整。

但这里有个容易被忽略的细节：双向扫描并不是简单地把序列反转再跑一次，而是需要维护两套独立的隐状态，最后在通道维度上做拼接或相加。工程实现时如果偷懒直接复用同一套状态，结果会明显掉点。我在复现类似结构时踩过这个坑，后面在实操章节会具体讲怎么排查。另外，ViM还讨论了跨扫描（cross-scan）的变体，把图像按不同方向切分再扫描，进一步提升空间连续性，但这会成倍增加计算量，属于精度和效率之间的取舍项，不是默认必选项。

3. 核心组件逐个拆：SSM、扫描策略与架构堆叠

3.1 状态空间模型的最小可用理解

抛开论文里那套连续时间微分方程，你只需要抓住离散化之后的三个矩阵：状态转移矩阵A、输入投影矩阵B、输出投影矩阵C。隐状态h的更新可以写成 h_t = A_bar * h_{t-1} + B_bar * x_t，输出 y_t = C * h_t。A_bar和B_bar是离散化之后的版本，离散化步长Δ决定了“每次更新记忆的幅度”。Δ大，记忆更新激进，适合快速变化的内容；Δ小，记忆保守，适合平滑延续的内容。

ViM的选择性机制就是让Δ、B、C都变成输入x_t的函数。这里有个实操要点：Δ通常通过一个线性层加softplus激活得到，要保证它是正数，否则离散化会失去物理意义。B和C则直接由线性投影生成。A矩阵一般不做输入依赖，保持为可学习参数，且常用对角结构以降低参数量和计算量。这套设定在多个开源实现里基本一致，你如果自己动手，建议先从对角A开始，别一上来就搞全矩阵，否则显存和调试难度都会陡增。

3.2 扫描顺序：ViM的“排队规则”设计

前面说过扫描顺序决定信息流向。ViM默认采用的是把二维patch序列按行优先拉平，然后做正反双向扫描。这个选择背后的逻辑是：行优先拉平至少保证了同一行内的相邻patch在序列里也是相邻的，水平方向的空间连续性得以保留；垂直方向的连续性则依赖双向扫描和后续层的堆叠来间接弥补。相比之下，如果按列优先或者Z字形拉平，效果会有差异，论文里也做过对比，行优先在多数分类任务上是稳妥的默认项。

但这里要提醒一句：扫描顺序的影响在浅层比较明显，到了深层，由于感受野已经通过多层传递扩展开来，顺序的敏感度会下降。所以如果你在做密集预测任务（比如分割），不要只盯着第一层的扫描方式调，更值得关注的是中间层是否引入了跨方向的信息交互。我见过一些复现把精力全花在改扫描顺序上，结果提升微乎其微，反而忽略了归一化和残差连接这些更基础的东西。

3.3 整体架构：ViM块长什么样

一个ViM块的结构可以粗略概括为：输入先过一层归一化，然后分成两条支路，一条走SSM扫描（内部包含前面说的选择性参数生成和双向扫描），另一条是残差直连；两条支路汇合后再过一层归一化和一个前馈网络（通常是MLP），最后再加一次残差。这个结构和Transformer块非常像，区别就是把多头自注意力换成了SSM扫描模块。这种“换芯不换壳”的设计有个实际好处：已有的训练框架、优化器配置、数据增强流程基本可以复用，迁移成本低。

堆叠方式上，ViM采用分层结构，类似金字塔，不同阶段做下采样，通道数逐级翻倍，序列长度逐级减半。这一点和Swin那类层次化Transformer思路一致，目的是兼顾分类任务需要的全局语义和密集预测需要的多尺度特征。下采样通常用简单的卷积或patch merging实现，具体选型对最终精度有影响，但不如SSM模块本身的影响大。我的建议是初期复现直接用patch merging，稳定之后再考虑替换成更复杂的下采样模块做消融。

4. 实操复现：从零搭一个可跑的ViM块

4.1 环境与依赖的务实选择

动手之前先把环境理清楚。ViM的核心扫描操作如果纯用PyTorch写循环，速度会慢到无法接受，所以实际实现通常依赖两类方案：一是调用经过编译优化的选择性扫描算子（很多开源实现会提供CUDA版本），二是用框架自带的高效扫描原语做近似。我的建议是，如果你只是想做算法验证和小规模实验，优先用纯PyTorch的参考实现，虽然慢但逻辑透明，方便调试；如果要跑大规模训练，再切换到优化过的算子，但要注意算子版本和框架版本的匹配，这里不匹配导致的报错往往很隐蔽。

依赖方面，基础的PyTorch、einops（做张量重排非常方便）、以及一个能跑通的扫描实现就够了。不建议一上来就引入一堆第三方库，依赖越多，出问题时定位越困难。我个人的习惯是先用最小依赖跑通一个block的前向，确认输出shape和数值范围正常，再逐步加东西。

4.2 关键参数的计算与选择过程

几个必须想清楚的参数：隐状态维度d_state、扩展因子expand、以及扫描的组数。d_state决定记忆的容量，太小记不住东西，太大显存吃不消，常见取值在16左右起步，具体要结合你的通道数调整。expand因子控制SSM内部通道的扩张倍数，一般取2，这个和Transformer里FFN的扩张思路类似，目的是给非线性变换留出空间。扫描组数则影响并行度和表达能力的平衡，组数多并行好但每组记忆容量小。

这里给一个粗略的估算方法：假设输入通道为C，expand取2，那么SSM内部通道约为2C，隐状态维度取d_state，则每个位置的隐状态参数量大致是2C乘以d_state量级。你可以据此反推显存占用，再决定是否要减小d_state或expand。我实测下来，在中等规模图像分类任务上，d_state从16降到8带来的精度损失通常在可接受范围内，但显存和速度改善明显，属于性价比很高的调参方向。

4.3 前向流程的逐步实现要点

第一步，输入张量形状通常是(batch, seq_len, channels)，先做归一化。第二步，通过线性层生成Δ、B、C，注意Δ要过softplus保证为正，并且通常会对Δ做一个范围裁剪防止数值爆炸。第三步，对A做离散化，得到A_bar，这一步涉及矩阵指数运算，实现时要注意数值稳定性，常见做法是对A取负保证其稳定。第四步，执行双向扫描，正向一遍反向一遍，各自得到输出后融合。第五步，过输出投影和残差，再进前馈网络。

每一步都有坑。比如Δ的裁剪阈值设多少，论文里给了参考值，但不同数据分布下最优值会漂移，建议把它做成可配置项而不是写死。再比如反向扫描时序列反转和状态初始化的顺序，写反了不会报错但结果会悄悄变差，这种bug最难查。我的做法是在小批量数据上同时跑正向和反向的手写循环版本，和向量化版本对比数值，确保两者一致后再上大规模训练。

4.4 训练配置的起步建议

优化器用AdamW，学习率从较小的值起步（比如1e-4量级），权重衰减正常设置。由于SSM的递归特性，梯度在长序列上可能出现不稳定，建议开启梯度裁剪，阈值设在1.0附近先观察。学习率调度用余弦退火比较稳妥。batch size受显存限制时，优先保证序列长度不要被过度压缩，因为ViM的优势恰恰在长序列上，你把序列砍太短就体现不出它的价值了。

还有一个容易被忽视的点：初始化。SSM相关参数的初始化对训练稳定性影响很大，A的初始化尤其关键，常见做法是让A的实部为负且分布在一个合理区间。如果你发现训练初期loss剧烈震荡，先别急着调学习率，回头检查一下A和Δ的初始化，十有八九是这里的问题。

5. 常见问题与排查技巧实录

5.1 训练不收敛或loss震荡

这是复现ViM时最高频的问题。排查顺序我一般是这样：先看Δ的取值范围，打印出来看看有没有异常大或异常小的值，softplus之后如果出现极端值，说明前面的线性层初始化有问题；再看A的离散化是否数值稳定，A_bar如果出现inf或nan，基本就是矩阵指数那步炸了；最后看梯度范数，如果梯度在少数层异常大，考虑加梯度裁剪或者调整该层的初始化。实测下来，大部分不收敛问题都能在前两步定位到。

5.2 显存占用远超预期

很多人以为线性复杂度就等于省显存，其实不然。ViM的隐状态需要在整个序列上维护，双向扫描还要维护两份，加上中间激活值，显存占用并不低。优化方向有几个：减小d_state、减小expand、用梯度检查点换显存、以及确保扫描算子是真正的高效实现而不是退化成循环。我遇到过一种情况，某实现因为版本问题悄悄退化成了逐元素循环，速度慢了几十倍但没有任何报错，排查了半天才发现。所以性能异常时，先确认你调用的算子真的在高效路径上。

5.3 精度上不去，分类指标卡住

如果训练正常但精度不理想，先做消融确认SSM模块是否真的在起作用。一个简单办法是把SSM输出置零只留残差，看精度掉多少，如果几乎不掉，说明模块没学到东西。常见原因包括：扫描方向单一导致空间信息利用不足、归一化位置放错、或者下采样太激进导致浅层信息丢失。另外，ViM对数据增强的敏感度和Transformer类似，增强太弱容易过拟合，太强又可能欠拟合，建议从中等强度起步做对比。

5.4 问题速查表

6. 影响范围与适用边界的一些个人判断

ViM这类模型真正有价值的地方，不在于它在某个分类榜单上刷了多少点，而在于它给“长序列视觉建模”提供了一条不同于注意力的技术路径。高分辨率图像、视频帧序列、遥感大图这些场景，序列长度动辄几千上万，注意力的平方开销是硬伤，而SSM的线性特性在这里有天然优势。我判断它在中长序列、对延迟敏感的部署场景里会更有生命力，而不是在短序列的小图上和成熟骨干硬拼。

但它也有明确的边界。选择性扫描的串行本质决定了它在某些硬件上的并行效率不如纯矩阵乘法，工程优化空间还很大；双向扫描带来的状态维护成本，让它在极长序列上的显存优势没有理论分析那么夸张。所以我的态度是：把它当成工具箱里的一把新扳手，适合拧特定型号的螺丝，而不是指望它替换掉所有现有工具。实际选型时，先明确你的序列长度、延迟预算和精度要求，再决定要不要上ViM，比盲目追新要靠谱得多。

最后分享一个我在复现过程中养成的小习惯：每加一个新模块，先用固定随机种子跑一个极小的过拟合实验，让模型去拟合几张图，如果能迅速把训练loss压到接近零，说明前向反向链路是通的，再去跑真实数据。这个土办法帮我省下了大量在“到底是模型不行还是代码有bug”之间反复横跳的时间，你也可以试试。