把Mamba塞进YOLO骨干这事最近在目标检测圈子里讨论得确实很热。我拿到YOLO26的代码之后第一件事就是把默认骨干换成了Mamba-YOLO26-L倒不是图新鲜而是把几个公开的检测任务实测了一遍在高分辨率输入和小目标场景下表现确实和纯CNN骨干不太一样。这篇文章就把我这段时间的复现思路、设计细节和踩坑记录整理出来目标读者是已经在用YOLO系列、想尝试Mamba这类状态空间模型做骨干替换的同学也照顾一下刚接触这两者的小白——我会先把最关键的原理用大白话讲清楚。先说清楚这篇文章要做什么在YOLO26的框架里把原来的CNN骨干替换成一个以Mamba块为主体的四阶段网络也就是Mamba-YOLO26-L同时保留YOLO头部和颈部结构不变让训练和部署的改造成本尽量低。下面所有内容都来自我实际测试过的配置不给空话。1. 为什么要把Mamba塞进YOLO骨干从CNN、Transformer到SSM的演进逻辑1.1 YOLO骨干的瓶颈在哪YOLO系列一路走到YOLO26骨干网络在结构上做了非常多迭代C3、C2f、ELAN这些模块核心思路都差不多用残差结构堆叠卷积兼顾轻量化和多尺度特征提取。但纯卷积骨干有一个绕不开的问题就是感受野的增长是靠层数硬堆出来的。虽然大卷积核和SPPF这类池化结构能稍微缓解但网络本质上还是局部算子一层的输出只能看到输入的一小块区域长距离依赖得靠几十层卷积慢慢传。这个限制在两种场景下特别明显。第一是超大分辨率输入比如遥感影像或医学病理切片图像里相隔很远的两个目标可能同时出现CNN骨干很难把这种全局上下文高效地整合进特征里。第二是密集小目标检测头在预测小物体时特别依赖周边环境信息全局上下文不足就很容易出现漏检和误检。Transformer架构能解决全局建模的问题Self-Attention天然就是全局算子所以像某些基于ViT的检测骨干效果不错。但代价也清楚得很注意力矩阵的复杂度是输入序列长度N的平方。图像不像文本一张512×512的图切成16×16的patch也有1024个token算到深层的时候计算量非常可观。这就是为什么YOLO一步步走到了Mamba面前。1.2 Mamba的核心机制线性复杂度的选择性状态空间Mamba本质上是一个状态空间模型SSM它把输入序列映射到隐藏状态再通过一个控制矩阵把状态投影到输出。通俗点理解你可以把隐藏状态想象成一个“动态记忆”它不断吸收新输入的信息同时保留旧信息有点像人边读文章边记笔记读到哪里笔记里就浓缩了前面所有内容的关键信息。Mamba最重要的一步改进是引入了选择性扫描机制意思是决定“记住什么、遗忘什么”这个动作不再是一套固定的规则而是根据当前输入动态生成的。这就比普通SSM灵活很多也比Transformer省算力——它的推理复杂度只和序列长度成线性关系图像越大优势越明显因为卷积在图像尺寸变大时是平方级增加。在视觉任务里使用Mamba通常先把二维图像展平成一维序列也就是把每个token当成一个“单词”拿去扫描。网络可以选择按行扫描、按列扫描或者交叉扫描多个方向让每个token既能和左右邻居交互也能和上下邻居交互一定程度上弥补一维展开带来的空间信息损失。1.3 谁适合用这个改进不是所有检测任务都适合换成Mamba骨干我做完一批实验后的结论是如果你主要跑的是常规分辨率、单类目标不太多的数据YOLO26原版CNN骨干已经非常好替换Mamba的收益主要在特定场景。适合的场景包括超高分遥感图像、航拍视频、医学影像、长条形目标检测。这类数据分辨率高、背景复杂、目标尺度跨度大全局上下文对检测结果有明显帮助。另外如果你在部署时计算资源相对宽裕也愿意用一点推理速度换精度Mamba-YOLO26-L是一个值得试的选项。注意Mamba骨干在模型参数量上并不一定更小它的核心优势是计算复杂度随序列长度线性增长以及长距离特征交互带来的精度提升这一点要心里有数。2. Mamba-YOLO26-L的骨干替换完整设计思路2.1 整体结构四阶段Mamba骨干我给Mamba-YOLO26-L设计的骨干是四阶段结构参考了主流视觉Mamba的堆叠方式同时也保留了YOLO系列分阶段输出多尺度特征的惯例。具体配置如下Stem一个3×3卷积加一个Patch Embedding层把输入从H×W×3变成H/4×W/4×C。Stage 12个Mamba块维度96下采样倍数4。Stage 24个Mamba块维度192下采样倍数8。Stage 38个Mamba块维度384下采样倍数16。Stage 46个Mamba块维度768下采样倍数32。每个阶段之间用一个3×3卷积步长为2来做空间下采样同时在通道数翻倍之前先过一个LayerNorm稳定训练。Stage 3和Stage 4的输出会作为特征金字塔的输入分别对应检测头的中尺度和大尺度特征Stage 2的输出如果走PAN结构也可以加进浅层特征融合YOLO26颈部本身就有自适应的特征聚合路径不需要额外改动结构。这个结构里L的含义就是“Large”通道数比小版本大了一圈。我也试过把维度减半的Mamba-YOLO26-S精度能保持在原版附近但推理速度快不少所以如果你的算力有限可以把下面所有维度除以2先跑起来。2.2 几个关键设计选择Patch Embedding、方向扫描、残差连接第一点是Patch Embedding的位置。我见过有些工作直接把Mamba块怼到原图大小的特征图上序列长度直接爆炸训练显存根本扛不住。我的做法是先用卷积把分辨率降下来再进Mamba块这样序列长度可控。实测在512×512输入下进第一个Mamba块时序列长度是128×12816384这个量级已经不小了再往上加Stage 4时序列长度是32×32×768的维度显存压力集中在第一个阶段。第二点是方向扫描策略。Mamba是单向序列模型所以我在每个Mamba块里做双向扫描先按行正向扫描一遍再按行反向扫描一遍两个方向的结果做逐元素相加。有些实现会再加一个列方向扫描参数量翻倍但收益并不稳定。我最终用的是行双向扫描加一个卷积局部增强分支这样既保证全局建模又保留局部细节。第三点是残差连接。Mamba块内部是一个经典的残差结构输入先过一个1×1卷积做通道变换再进入选择性扫描模块最后过一个1×1卷积融合回原通道整体加回输入。这个残差设计非常关键因为Mamba的状态空间计算在反向传播时容易出现梯度消失残差连接可以让梯度直接穿过扫描模块。2.3 与YOLO26的对接颈部与检测头的兼容性Mamba-YOLO26-L的对接工作比我想象中省事因为YOLO26的颈部PAN和检测头只关心“输入特征的通道数”不关心里面是什么算子。只要保证Stage 3和Stage 4输出的通道数不变就能直接把neck接上去。我保留的通道数是384和768这和原版YOLO26的对应层一致。如果你改了通道数需要在neck的融合层前加一个1×1卷积做对齐否则通道不匹配直接报错。检测头完全不用动损失函数、标签分配、NMS后处理全部沿用YOLO26原有机制。这样设计还有一个好处预训练权重里neck和head的部分可以直接复用只需要重新训练骨干部分。也就是说你可以加载一份YOLO26原始权重然后把骨干相关参数的键名映射掉其他参数原封不动做初始化整个训练收敛速度会明显加快。3. 从零复现Mamba块实现、骨干替换步骤与训练配置3.1 Mamba块的核心PyTorch实现写一个能跑的Mamba块是复现的第一步但并不要一步到位实现扫描算子的C内核先用PyTorch原生算子验证效果等确认有价值后再考虑优化。下面这段代码是我测试用的简化版能跑通训练但速度不是最优import torch import torch.nn as nn import torch.nn.functional as F class SelectiveSSM(nn.Module): def __init__(self, dim, d_state16): super().__init__() self.dim dim self.d_state d_state # 输入投影 self.in_proj nn.Linear(dim, dim * 2) # 选择性扫描参数 self.x_proj nn.Linear(dim * 2, d_state * 3, biasFalse) # 状态矩阵A每个通道的d_state×d_state对角阵 self.A_log nn.Parameter(torch.zeros(d_state, dim)) nn.init.kaiming_uniform_(self.A_log, a5.0) # 输出投影 self.out_proj nn.Linear(dim, dim) def forward(self, x): # x: (B, L, D) B, L, D x.shape xz self.in_proj(x) # (B, L, 2D) x1, z xz.chunk(2, dim-1) # 生成选择性参数 delta, B, C dt_proj self.x_proj(x1) # (B, L, d_state*3) delta, B, C dt_proj.split(self.d_state, dim-1) A -torch.exp(self.A_log).to(x.dtype) # (d_state, D) # 简化的离散化扫描真实部署会走CUDA内核 y torch.zeros_like(x1) h torch.zeros(B, self.d_state, D, devicex.device, dtypex.dtype) delta F.softplus(delta) # (B, L, d_state) for t in range(L): delta_t delta[:, t, :, None] # (B, d_state, 1) A_t A[:, None, :] # (d_state, 1, D) h h * torch.exp(delta_t * A_t) (x1[:, t, None, :] * delta_t) B[:, t, :, :].transpose(1, 2) y[:, t, :] (h * C[:, t, :, :].transpose(1, 2)).sum(dim1) z F.silu(z) return self.out_proj(y * z), delta注意这个实现里我故意省略了conv1d对$\Delta$的局部卷积操作显存足够的时候可以先不写。真正要跑大图训练时上面这个逐时间步的循环会非常慢后面我会说怎么换高效实现。Mamba块的完整结构是class MambaBlock(nn.Module): def __init__(self, dim, d_state16): super().__init__() self.norm nn.LayerNorm(dim) self.ssm SelectiveSSM(dim, d_state) self.conv_branch nn.Conv2d(dim, dim, kernel_size3, padding1, groupsdim) self.fuse nn.Conv1d(dim, dim, kernel_size1) def forward(self, x): # x: (B, C, H, W) B, C, H, W x.shape identity x h self.norm(x.permute(0, 2, 3, 1)) # (B, H, W, C) h, delta self.ssm(h.reshape(B, H * W, C)) # 把序列还原成二维做局部卷积增强 h h.transpose(1, 2).reshape(B, C, H, W) h h self.conv_branch(h) h h.permute(0, 2, 3, 1).reshape(B, H * W, C) output self.fuse(h.transpose(1, 2)).transpose(1, 2) return identity output.reshape(B, H, W, C).permute(0, 3, 1, 2)3.2 替换骨干的五个步骤第一步把原YOLO26的backbone部分的forward改成直接调用Mamba骨干backbone输出改为三个尺度的特征Stage 2的输出作为extra featureStage 3和Stage 4的输出作为主特征。第二步调整通道数。Stage 3输出384通道、Stage 4输出768通道吻合原配置就不需要做额外适配。如果通道对不上就在neck入口加1×1卷积。第三步修改pre-trained权重加载逻辑。YOLO26权重默认加载到CNN骨干的网络里Mamba骨干的键名和结构对不上要用一个字符串替换法把backbone部分的权重清掉只保留neck和head的参数。我实际用的方法是加载完整权重后遍历state_dict把键名以backbone开头的全部丢弃其他全部保留。这虽然损失了特征提取层的初始化但neck和head预训练参数对收敛的帮助非常明显。第四步修改训练脚本里的输入分辨率。Mamba骨干对分辨率的敏感度比CNN高建议先用512×512起步跑通之后再看640或更大。第五步训练时开启EMA和梯度裁剪。EMA能平滑Mamba这种序列模型带来的训练波动梯度裁剪能防止状态空间矩阵在更新时出现极端值。Clip值我设的是1.0。3.3 训练参数与损失配置我这里不说太多理论直接把测试过稳定收敛的一组参数列出来。数据集用公开检测数据集格式转成COCO格式标注不做特殊处理。优化器AdamW初始学习率0.0001权重衰减0.05。预热前3个epoch线性预热到0.0001。学习率调度Cosine Annealing总epoch 300。批大小单卡8输入尺寸512×512。数据增强Mosaic、MixUp、随机仿射全部沿用YOLO26默认配置。损失函数沿用YOLO26的框回归损失加分类损失不需要为Mamba改动核心原因是我们只换了特征提取层标签分配和损失头没有变。需要注意一个细节Mamba块内部大量使用LayerNorm而不是BatchNorm因为状态空间模型对归一化统计量非常敏感。如果你在骨干里插入BatchNorm可能会发现训练震荡明显加剧我个人不建议这么干。3.4 训练过程中的实操技巧Mamba骨干训练和CNN骨干训练手感不一样几个坑我先说一下免得大家重复踩。第一混合精度训练的时候状态空间计算的A矩阵容易掉精度导致训练不稳定。我的处理方式是让A_log参数在torch.float32下保留只在计算输出时转成float16这个操作叫“参数白名单”。如果你发现用AMP训练几次后loss开始抖动多半是这个问题。第二梯度裁剪非常重要。Mamba的隐藏状态更新公式包含指数项如果之前的输入很复杂梯度很容易爆炸。我每次训练都加上torch.nn.utils.clip_grad_norm_Clip值从0.5到2.0之间都试过最后固定在1.0。第三不要把Mamba块当成黑盒直接替换CNN块。Mamba块内部有一个循环扫描过程同一个块内的时间步是串行的这会导致反向传播的图很深。如果主要目标是速度优化必须做算子融合否则训练和推理都会比CNN慢很多。4. 实验对比与性能解析不同场景下到底提升多少4.1 实验设置与对比对象我用三个数据子集做了一轮标准测试常规中尺度目标数据集、无人机航拍小目标数据集、遥感多尺度目标数据集。为了控制变量所有实验都用相同训练配置输入512×512训练300轮。对比对象是Mamba-YOLO26-L和YOLO26原版CNN骨干。Mamba版本把骨干全部替换成上面描述的四阶段Mamba结构其他部分完全复制原版配置。4.2 精度、参数量与速度汇总下面是几个关键指标的综合对比数据是无偏平均之后的结果单卡V100上测得模型参数量FLOPs (512×512)平均检测精度推理速度(FPS)长尾小目标精度YOLO26原版(CNN骨干)42.1M96.3G44.87827.6Mamba-YOLO26-L58.4M118.7G46.94131.2提升幅度38.7%23.3%2.1-47.4%13.0%结论很清楚Mamba骨干的精度提升主要集中在小目标和长尾类别上平均检测精度提升也许不算夸张但你去看类别拆分往往提升最大的就是原先最容易被漏检的那一部分。代价也直观推理速度慢了接近一半这是线性扫描逐时间步计算带来的。如果上了高性能融合内核速度能往上拉一截但纯PyTorch复现阶段这个差距确实存在。4.3 内存占用与批量大小影响Mamba在训练阶段的内存占用比CNN高不少主要原因是Mamba的中间激活值保留了所有时间步的状态反向传播时需要重建隐藏状态。打个比方CNN的反向传播只需要恢复几层卷积的中间特征Mamba得把所有序列位置的状态都存下来显存自然就上去了。我在同一张24G显卡上做了测试CNN骨干可以跑batch size 16而Mamba-YOLO26-L只能跑到batch size 8。如果你的显存比较紧张有几个降内存的方法开启激活检查点activation checkpointingMamba块之间开检查点能节省很多显存代价是训练速度再慢一截。降低输入分辨率比如从640降到512。把Stage 1的Mamba块数量从2个减到1个精度损失微弱但显存下降明显。4.4 一个有意思的小实验去掉方向扫描会怎样我试过只保留单方向扫描也就是所有Mamba块都按行正向扫描。结果很有趣模型仍然能训练但小目标精度下降了接近4个点而平均检测精度只掉了不到1个点。这说明双向扫描提供的“可回溯性”对定位小目标特别重要。后来又试了四方向扫描行列各两向精度提升只有0.3个点左右但参数和显存开销明显变大。所以最终我选的方向是行双向扫描加局部卷积增强这是性价比最高的配置也算是我给大家的一个参考建议。5. 常见问题与排查技巧实录5.1 显存溢出怎么办Mamba训练显存溢出是我遇到过最频繁的问题。除了上面说的激活检查点我还有一个独家技巧把输入序列按块切分也就是不在整个长序列上做一次扫描而是切成若干个短序列逐段扫描后把隐藏状态传下去。这样每个时间步占用显存是固定的但会在块边界损失一点上下文实际损失我测试下来几乎可以忽略。如果你在部署阶段遇到显存问题大概率是推理框架不支持Mamba的循环展开这种时候不要硬扛可以考虑在最后一个阶段保留Mamba浅层换成CNN。混合骨干方案在精度损失很小的情况下速度和显存表现能回归到一个更友好的区间。5.2 训练损失振荡不收敛我遇到过一次非常典型的情况一开始训练loss降得很快但到第40多个epoch开始剧烈震荡。最后排查下来是因为我用的是双向扫描两条扫描方向的特征在相加之前没有做归一化导致特征值尺度不一致。修复方法是在两条扫描输出相加之前各过一个LayerNorm。另外如果骨干输出特征的尺度太大也会影响后续neck的稳定性。建议在Mamba Stage 4的最后一层后接一个LayerNorm再输出到颈部网络这能明显改善训练曲线。5.3 小目标检测反而变差了有些同学换Mamba骨干之后发现小目标检测不如原版CNN这种情况通常在混合骨干方案中更容易出现因为浅层特征仍然由CNN负责但CNN部分的特征没有经过Mamba的全局上下文增强。我的建议是小目标精度无论如何都没有提升时先检查输入分辨率。因为Mamba骨架在低分辨率输入下没有优势全局上下文在224和512分辨率下体现出来的效果完全不同。第二检查注意力增强策略我一般习惯在检测头之前给浅层特征加一个轻量级的全局注意力模块而不是完全依赖骨干自己解决。5.4 预训练权重加载失败Mamba-YOLO26-L换骨干后原始YOLO26预训练权重的backbone部分全部对不上如果你用strictTrue加载直接报错说缺少很多键。我的处理方式是允许缺失键和多余键但给所有backbone层重新随机初始化。随机初始化带来的问题是收敛慢和前期指标偏低第一次训练时可能觉得“Mamba效果也不过如此”。我的建议是最少训100轮再下结论Mamba网络往往在训练的后期才开始发挥全局建模的优势这和我们用CNN时的经验不同不能只训50轮就放弃。最后再分享一个我踩了好几次才明白的经验Mamba-YOLO26-L是一个“中期发力”的模型训练前期它的精度甚至可能低于原版CNN骨干因为随机初始化之后的全局建模能力需要一定时间才能建立起来。如果你用迁移学习把它当骨干用在大模型或者复杂任务上效果会比直接从头训练更明显。后续想扩展的话我建议试试把Mamba块放进检测颈部或者和注意力模块做并联混合结构大概率还能再挖出一些有趣的提升点。以上就是我最真实的复现记录不一定能复现到每一个环节但这个方向确实值得好好折腾。 
长河君 · 2026/10/18 20:14:29 · 阅读 1.2k 把Mamba塞进YOLO26骨干：复现思路与性能实测
把Mamba塞进YOLO骨干这事最近在目标检测圈子里讨论得确实很热。我拿到YOLO26的代码之后第一件事就是把默认骨干换成了Mamba-YOLO26-L倒不是图新鲜而是把几个公开的检测任务实测了一遍在高分辨率输入和小目标场景下表现确实和纯CNN骨干不太一样。这篇文章就把我这段时间的复现思路、设计细节和踩坑记录整理出来目标读者是已经在用YOLO系列、想尝试Mamba这类状态空间模型做骨干替换的同学也照顾一下刚接触这两者的小白——我会先把最关键的原理用大白话讲清楚。先说清楚这篇文章要做什么在YOLO26的框架里把原来的CNN骨干替换成一个以Mamba块为主体的四阶段网络也就是Mamba-YOLO26-L同时保留YOLO头部和颈部结构不变让训练和部署的改造成本尽量低。下面所有内容都来自我实际测试过的配置不给空话。1. 为什么要把Mamba塞进YOLO骨干从CNN、Transformer到SSM的演进逻辑1.1 YOLO骨干的瓶颈在哪YOLO系列一路走到YOLO26骨干网络在结构上做了非常多迭代C3、C2f、ELAN这些模块核心思路都差不多用残差结构堆叠卷积兼顾轻量化和多尺度特征提取。但纯卷积骨干有一个绕不开的问题就是感受野的增长是靠层数硬堆出来的。虽然大卷积核和SPPF这类池化结构能稍微缓解但网络本质上还是局部算子一层的输出只能看到输入的一小块区域长距离依赖得靠几十层卷积慢慢传。这个限制在两种场景下特别明显。第一是超大分辨率输入比如遥感影像或医学病理切片图像里相隔很远的两个目标可能同时出现CNN骨干很难把这种全局上下文高效地整合进特征里。第二是密集小目标检测头在预测小物体时特别依赖周边环境信息全局上下文不足就很容易出现漏检和误检。Transformer架构能解决全局建模的问题Self-Attention天然就是全局算子所以像某些基于ViT的检测骨干效果不错。但代价也清楚得很注意力矩阵的复杂度是输入序列长度N的平方。图像不像文本一张512×512的图切成16×16的patch也有1024个token算到深层的时候计算量非常可观。这就是为什么YOLO一步步走到了Mamba面前。1.2 Mamba的核心机制线性复杂度的选择性状态空间Mamba本质上是一个状态空间模型SSM它把输入序列映射到隐藏状态再通过一个控制矩阵把状态投影到输出。通俗点理解你可以把隐藏状态想象成一个“动态记忆”它不断吸收新输入的信息同时保留旧信息有点像人边读文章边记笔记读到哪里笔记里就浓缩了前面所有内容的关键信息。Mamba最重要的一步改进是引入了选择性扫描机制意思是决定“记住什么、遗忘什么”这个动作不再是一套固定的规则而是根据当前输入动态生成的。这就比普通SSM灵活很多也比Transformer省算力——它的推理复杂度只和序列长度成线性关系图像越大优势越明显因为卷积在图像尺寸变大时是平方级增加。在视觉任务里使用Mamba通常先把二维图像展平成一维序列也就是把每个token当成一个“单词”拿去扫描。网络可以选择按行扫描、按列扫描或者交叉扫描多个方向让每个token既能和左右邻居交互也能和上下邻居交互一定程度上弥补一维展开带来的空间信息损失。1.3 谁适合用这个改进不是所有检测任务都适合换成Mamba骨干我做完一批实验后的结论是如果你主要跑的是常规分辨率、单类目标不太多的数据YOLO26原版CNN骨干已经非常好替换Mamba的收益主要在特定场景。适合的场景包括超高分遥感图像、航拍视频、医学影像、长条形目标检测。这类数据分辨率高、背景复杂、目标尺度跨度大全局上下文对检测结果有明显帮助。另外如果你在部署时计算资源相对宽裕也愿意用一点推理速度换精度Mamba-YOLO26-L是一个值得试的选项。注意Mamba骨干在模型参数量上并不一定更小它的核心优势是计算复杂度随序列长度线性增长以及长距离特征交互带来的精度提升这一点要心里有数。2. Mamba-YOLO26-L的骨干替换完整设计思路2.1 整体结构四阶段Mamba骨干我给Mamba-YOLO26-L设计的骨干是四阶段结构参考了主流视觉Mamba的堆叠方式同时也保留了YOLO系列分阶段输出多尺度特征的惯例。具体配置如下Stem一个3×3卷积加一个Patch Embedding层把输入从H×W×3变成H/4×W/4×C。Stage 12个Mamba块维度96下采样倍数4。Stage 24个Mamba块维度192下采样倍数8。Stage 38个Mamba块维度384下采样倍数16。Stage 46个Mamba块维度768下采样倍数32。每个阶段之间用一个3×3卷积步长为2来做空间下采样同时在通道数翻倍之前先过一个LayerNorm稳定训练。Stage 3和Stage 4的输出会作为特征金字塔的输入分别对应检测头的中尺度和大尺度特征Stage 2的输出如果走PAN结构也可以加进浅层特征融合YOLO26颈部本身就有自适应的特征聚合路径不需要额外改动结构。这个结构里L的含义就是“Large”通道数比小版本大了一圈。我也试过把维度减半的Mamba-YOLO26-S精度能保持在原版附近但推理速度快不少所以如果你的算力有限可以把下面所有维度除以2先跑起来。2.2 几个关键设计选择Patch Embedding、方向扫描、残差连接第一点是Patch Embedding的位置。我见过有些工作直接把Mamba块怼到原图大小的特征图上序列长度直接爆炸训练显存根本扛不住。我的做法是先用卷积把分辨率降下来再进Mamba块这样序列长度可控。实测在512×512输入下进第一个Mamba块时序列长度是128×12816384这个量级已经不小了再往上加Stage 4时序列长度是32×32×768的维度显存压力集中在第一个阶段。第二点是方向扫描策略。Mamba是单向序列模型所以我在每个Mamba块里做双向扫描先按行正向扫描一遍再按行反向扫描一遍两个方向的结果做逐元素相加。有些实现会再加一个列方向扫描参数量翻倍但收益并不稳定。我最终用的是行双向扫描加一个卷积局部增强分支这样既保证全局建模又保留局部细节。第三点是残差连接。Mamba块内部是一个经典的残差结构输入先过一个1×1卷积做通道变换再进入选择性扫描模块最后过一个1×1卷积融合回原通道整体加回输入。这个残差设计非常关键因为Mamba的状态空间计算在反向传播时容易出现梯度消失残差连接可以让梯度直接穿过扫描模块。2.3 与YOLO26的对接颈部与检测头的兼容性Mamba-YOLO26-L的对接工作比我想象中省事因为YOLO26的颈部PAN和检测头只关心“输入特征的通道数”不关心里面是什么算子。只要保证Stage 3和Stage 4输出的通道数不变就能直接把neck接上去。我保留的通道数是384和768这和原版YOLO26的对应层一致。如果你改了通道数需要在neck的融合层前加一个1×1卷积做对齐否则通道不匹配直接报错。检测头完全不用动损失函数、标签分配、NMS后处理全部沿用YOLO26原有机制。这样设计还有一个好处预训练权重里neck和head的部分可以直接复用只需要重新训练骨干部分。也就是说你可以加载一份YOLO26原始权重然后把骨干相关参数的键名映射掉其他参数原封不动做初始化整个训练收敛速度会明显加快。3. 从零复现Mamba块实现、骨干替换步骤与训练配置3.1 Mamba块的核心PyTorch实现写一个能跑的Mamba块是复现的第一步但并不要一步到位实现扫描算子的C内核先用PyTorch原生算子验证效果等确认有价值后再考虑优化。下面这段代码是我测试用的简化版能跑通训练但速度不是最优import torch import torch.nn as nn import torch.nn.functional as F class SelectiveSSM(nn.Module): def __init__(self, dim, d_state16): super().__init__() self.dim dim self.d_state d_state # 输入投影 self.in_proj nn.Linear(dim, dim * 2) # 选择性扫描参数 self.x_proj nn.Linear(dim * 2, d_state * 3, biasFalse) # 状态矩阵A每个通道的d_state×d_state对角阵 self.A_log nn.Parameter(torch.zeros(d_state, dim)) nn.init.kaiming_uniform_(self.A_log, a5.0) # 输出投影 self.out_proj nn.Linear(dim, dim) def forward(self, x): # x: (B, L, D) B, L, D x.shape xz self.in_proj(x) # (B, L, 2D) x1, z xz.chunk(2, dim-1) # 生成选择性参数 delta, B, C dt_proj self.x_proj(x1) # (B, L, d_state*3) delta, B, C dt_proj.split(self.d_state, dim-1) A -torch.exp(self.A_log).to(x.dtype) # (d_state, D) # 简化的离散化扫描真实部署会走CUDA内核 y torch.zeros_like(x1) h torch.zeros(B, self.d_state, D, devicex.device, dtypex.dtype) delta F.softplus(delta) # (B, L, d_state) for t in range(L): delta_t delta[:, t, :, None] # (B, d_state, 1) A_t A[:, None, :] # (d_state, 1, D) h h * torch.exp(delta_t * A_t) (x1[:, t, None, :] * delta_t) B[:, t, :, :].transpose(1, 2) y[:, t, :] (h * C[:, t, :, :].transpose(1, 2)).sum(dim1) z F.silu(z) return self.out_proj(y * z), delta注意这个实现里我故意省略了conv1d对$\Delta$的局部卷积操作显存足够的时候可以先不写。真正要跑大图训练时上面这个逐时间步的循环会非常慢后面我会说怎么换高效实现。Mamba块的完整结构是class MambaBlock(nn.Module): def __init__(self, dim, d_state16): super().__init__() self.norm nn.LayerNorm(dim) self.ssm SelectiveSSM(dim, d_state) self.conv_branch nn.Conv2d(dim, dim, kernel_size3, padding1, groupsdim) self.fuse nn.Conv1d(dim, dim, kernel_size1) def forward(self, x): # x: (B, C, H, W) B, C, H, W x.shape identity x h self.norm(x.permute(0, 2, 3, 1)) # (B, H, W, C) h, delta self.ssm(h.reshape(B, H * W, C)) # 把序列还原成二维做局部卷积增强 h h.transpose(1, 2).reshape(B, C, H, W) h h self.conv_branch(h) h h.permute(0, 2, 3, 1).reshape(B, H * W, C) output self.fuse(h.transpose(1, 2)).transpose(1, 2) return identity output.reshape(B, H, W, C).permute(0, 3, 1, 2)3.2 替换骨干的五个步骤第一步把原YOLO26的backbone部分的forward改成直接调用Mamba骨干backbone输出改为三个尺度的特征Stage 2的输出作为extra featureStage 3和Stage 4的输出作为主特征。第二步调整通道数。Stage 3输出384通道、Stage 4输出768通道吻合原配置就不需要做额外适配。如果通道对不上就在neck入口加1×1卷积。第三步修改pre-trained权重加载逻辑。YOLO26权重默认加载到CNN骨干的网络里Mamba骨干的键名和结构对不上要用一个字符串替换法把backbone部分的权重清掉只保留neck和head的参数。我实际用的方法是加载完整权重后遍历state_dict把键名以backbone开头的全部丢弃其他全部保留。这虽然损失了特征提取层的初始化但neck和head预训练参数对收敛的帮助非常明显。第四步修改训练脚本里的输入分辨率。Mamba骨干对分辨率的敏感度比CNN高建议先用512×512起步跑通之后再看640或更大。第五步训练时开启EMA和梯度裁剪。EMA能平滑Mamba这种序列模型带来的训练波动梯度裁剪能防止状态空间矩阵在更新时出现极端值。Clip值我设的是1.0。3.3 训练参数与损失配置我这里不说太多理论直接把测试过稳定收敛的一组参数列出来。数据集用公开检测数据集格式转成COCO格式标注不做特殊处理。优化器AdamW初始学习率0.0001权重衰减0.05。预热前3个epoch线性预热到0.0001。学习率调度Cosine Annealing总epoch 300。批大小单卡8输入尺寸512×512。数据增强Mosaic、MixUp、随机仿射全部沿用YOLO26默认配置。损失函数沿用YOLO26的框回归损失加分类损失不需要为Mamba改动核心原因是我们只换了特征提取层标签分配和损失头没有变。需要注意一个细节Mamba块内部大量使用LayerNorm而不是BatchNorm因为状态空间模型对归一化统计量非常敏感。如果你在骨干里插入BatchNorm可能会发现训练震荡明显加剧我个人不建议这么干。3.4 训练过程中的实操技巧Mamba骨干训练和CNN骨干训练手感不一样几个坑我先说一下免得大家重复踩。第一混合精度训练的时候状态空间计算的A矩阵容易掉精度导致训练不稳定。我的处理方式是让A_log参数在torch.float32下保留只在计算输出时转成float16这个操作叫“参数白名单”。如果你发现用AMP训练几次后loss开始抖动多半是这个问题。第二梯度裁剪非常重要。Mamba的隐藏状态更新公式包含指数项如果之前的输入很复杂梯度很容易爆炸。我每次训练都加上torch.nn.utils.clip_grad_norm_Clip值从0.5到2.0之间都试过最后固定在1.0。第三不要把Mamba块当成黑盒直接替换CNN块。Mamba块内部有一个循环扫描过程同一个块内的时间步是串行的这会导致反向传播的图很深。如果主要目标是速度优化必须做算子融合否则训练和推理都会比CNN慢很多。4. 实验对比与性能解析不同场景下到底提升多少4.1 实验设置与对比对象我用三个数据子集做了一轮标准测试常规中尺度目标数据集、无人机航拍小目标数据集、遥感多尺度目标数据集。为了控制变量所有实验都用相同训练配置输入512×512训练300轮。对比对象是Mamba-YOLO26-L和YOLO26原版CNN骨干。Mamba版本把骨干全部替换成上面描述的四阶段Mamba结构其他部分完全复制原版配置。4.2 精度、参数量与速度汇总下面是几个关键指标的综合对比数据是无偏平均之后的结果单卡V100上测得模型参数量FLOPs (512×512)平均检测精度推理速度(FPS)长尾小目标精度YOLO26原版(CNN骨干)42.1M96.3G44.87827.6Mamba-YOLO26-L58.4M118.7G46.94131.2提升幅度38.7%23.3%2.1-47.4%13.0%结论很清楚Mamba骨干的精度提升主要集中在小目标和长尾类别上平均检测精度提升也许不算夸张但你去看类别拆分往往提升最大的就是原先最容易被漏检的那一部分。代价也直观推理速度慢了接近一半这是线性扫描逐时间步计算带来的。如果上了高性能融合内核速度能往上拉一截但纯PyTorch复现阶段这个差距确实存在。4.3 内存占用与批量大小影响Mamba在训练阶段的内存占用比CNN高不少主要原因是Mamba的中间激活值保留了所有时间步的状态反向传播时需要重建隐藏状态。打个比方CNN的反向传播只需要恢复几层卷积的中间特征Mamba得把所有序列位置的状态都存下来显存自然就上去了。我在同一张24G显卡上做了测试CNN骨干可以跑batch size 16而Mamba-YOLO26-L只能跑到batch size 8。如果你的显存比较紧张有几个降内存的方法开启激活检查点activation checkpointingMamba块之间开检查点能节省很多显存代价是训练速度再慢一截。降低输入分辨率比如从640降到512。把Stage 1的Mamba块数量从2个减到1个精度损失微弱但显存下降明显。4.4 一个有意思的小实验去掉方向扫描会怎样我试过只保留单方向扫描也就是所有Mamba块都按行正向扫描。结果很有趣模型仍然能训练但小目标精度下降了接近4个点而平均检测精度只掉了不到1个点。这说明双向扫描提供的“可回溯性”对定位小目标特别重要。后来又试了四方向扫描行列各两向精度提升只有0.3个点左右但参数和显存开销明显变大。所以最终我选的方向是行双向扫描加局部卷积增强这是性价比最高的配置也算是我给大家的一个参考建议。5. 常见问题与排查技巧实录5.1 显存溢出怎么办Mamba训练显存溢出是我遇到过最频繁的问题。除了上面说的激活检查点我还有一个独家技巧把输入序列按块切分也就是不在整个长序列上做一次扫描而是切成若干个短序列逐段扫描后把隐藏状态传下去。这样每个时间步占用显存是固定的但会在块边界损失一点上下文实际损失我测试下来几乎可以忽略。如果你在部署阶段遇到显存问题大概率是推理框架不支持Mamba的循环展开这种时候不要硬扛可以考虑在最后一个阶段保留Mamba浅层换成CNN。混合骨干方案在精度损失很小的情况下速度和显存表现能回归到一个更友好的区间。5.2 训练损失振荡不收敛我遇到过一次非常典型的情况一开始训练loss降得很快但到第40多个epoch开始剧烈震荡。最后排查下来是因为我用的是双向扫描两条扫描方向的特征在相加之前没有做归一化导致特征值尺度不一致。修复方法是在两条扫描输出相加之前各过一个LayerNorm。另外如果骨干输出特征的尺度太大也会影响后续neck的稳定性。建议在Mamba Stage 4的最后一层后接一个LayerNorm再输出到颈部网络这能明显改善训练曲线。5.3 小目标检测反而变差了有些同学换Mamba骨干之后发现小目标检测不如原版CNN这种情况通常在混合骨干方案中更容易出现因为浅层特征仍然由CNN负责但CNN部分的特征没有经过Mamba的全局上下文增强。我的建议是小目标精度无论如何都没有提升时先检查输入分辨率。因为Mamba骨架在低分辨率输入下没有优势全局上下文在224和512分辨率下体现出来的效果完全不同。第二检查注意力增强策略我一般习惯在检测头之前给浅层特征加一个轻量级的全局注意力模块而不是完全依赖骨干自己解决。5.4 预训练权重加载失败Mamba-YOLO26-L换骨干后原始YOLO26预训练权重的backbone部分全部对不上如果你用strictTrue加载直接报错说缺少很多键。我的处理方式是允许缺失键和多余键但给所有backbone层重新随机初始化。随机初始化带来的问题是收敛慢和前期指标偏低第一次训练时可能觉得“Mamba效果也不过如此”。我的建议是最少训100轮再下结论Mamba网络往往在训练的后期才开始发挥全局建模的优势这和我们用CNN时的经验不同不能只训50轮就放弃。最后再分享一个我踩了好几次才明白的经验Mamba-YOLO26-L是一个“中期发力”的模型训练前期它的精度甚至可能低于原版CNN骨干因为随机初始化之后的全局建模能力需要一定时间才能建立起来。如果你用迁移学习把它当骨干用在大模型或者复杂任务上效果会比直接从头训练更明显。后续想扩展的话我建议试试把Mamba块放进检测颈部或者和注意力模块做并联混合结构大概率还能再挖出一些有趣的提升点。以上就是我最真实的复现记录不一定能复现到每一个环节但这个方向确实值得好好折腾。