最近在做目标检测模型的精度优化,一直在折腾YOLOv12的Neck部分。说实话,Neck这块一直是“有活干但不出彩”的位置,标准的PAN-FPN结构用到现在,融合是够用了,但特征精炼这件事一直没被认真对待。直到我看到arXiv 2025上那篇关于分层多头注意力HMHA的工作,思路是把多头注意力的头拆成多个层次子空间,再配合通道重排来降低冗余。我把这套改进接到了YOLOv12的Neck上,试跑了几轮实验,效果比我预想中要稳,尤其是小目标和重叠目标的召回率提升很直观。这篇文章就围绕这个改进,把HMHA的原理、代码实现、插入Neck的位置选择、参数调节顺序和踩过的坑完整记录下来,给正在做YOLOv12改进的同行一个可参考的实操版本。
1. 为什么Neck是注意力改进的主战场
1.1 YOLOv12 Neck的老问题:融合有余,精炼不足
YOLOv12的Neck结构并没有完全跳出经典的双向特征金字塔框架,自顶向下传递高层语义,自底向上补充底层细节,最后在P3、P4、P5三个尺度上输出融合特征。这套结构的核心操作就是concat和add,本质上是在做“特征搬运”,把不同尺度的张量拼接在一起。问题在于,concat操作会直接把通道翻倍,而相邻层级的语义差异又很大,比如P5感受野大、偏向全局语义,P3分辨率高、保留的是边缘和纹理,这两类特征硬拼在一起,必然产生大量冗余通道。
这种冗余带来的后果是:Neck里的卷积层需要花大量参数去“消化”那些重复表达,而不是专注于提取互补信息。我做过一个简单实验,把YOLOv12在Backbone输出的特征直接做一次通道相关性统计,发现P3和P4在部分通道上的相关系数能到0.7以上,也就是说相当一部分通道在重复表达相近的模式。通道一旦冗余,模型的表达能力就会下降,而参数量还在白白增加。这也是为什么很多人往Neck里加注意力模块,但效果忽高忽低的原因——注意力组件如果不去处理通道冗余,只是换个方式加权,本质上还是在“搬运特征”。
1.2 直接在Neck里加SE或普通MHSA为什么效果不稳定
不少同学会直接往Neck里加SE模块或者普通的多头自注意力MHSA。SE模块的思路是通道维度上做全局池化,然后再加权,操作轻量、速度影响小,但它只做了通道维度的重标定,完全没有考虑空间位置上哪些区域需要被强调。在目标检测里,同一张特征图上的不同位置重要性差异极大,SE这种“一把尺子量所有空间位置”的做法,在小目标聚集的场景里经常帮倒忙。
普通MHSA则是另一个极端,它确实能建模长距离依赖,但代价也很明显。首先是计算量爆炸,特征图分辨率稍高一点,注意力矩阵的尺寸就变成N×N,显存和时间都扛不住。其次是训练不稳定,我在实验里发现普通多头注意力的多个头在收敛后,注意力分布会趋向一致,经常出现三四个头的热力图几乎一样的情况,这种情况学到的特征多样性非常有限,叠加到Neck的融合特征上,不但没有补充新信息,反而放大了原有冗余。这其实就是因为MHSA的各个头共享同一套qkv投影初始化,如果没有外部约束,最终会收敛到相似的解。HMHA的思路恰恰是解决这两个问题:通过子空间拆分给不同头不同的尺度约束,通过通道重排打破头之间的冗余关联。
2. HMHA到底改了什么
2.1 分层多头注意力的核心思想
HMHA的全称是Hierarchical Multi-Head Attention,直译过来就是分层多头注意力。它的核心做法是把原本“平级”的注意力头分成若干组,每一组负责一个特征子空间。比如8个注意力头分成4组,每组2个头,四组分别在不同窗口范围或不同空间尺度下计算注意力。这个设计和人的视觉注意力机制有点相似,人们看一张图时,既会关注整体场景,也会聚焦到局部细节,这两件事是并行发生的。HMHA就是让一部分头负责全局建模,让另一部分头负责局部细节,各司其职。
分层带来的直接好处是注意力多样性提升。不同组的头天然被约束在不同的感知范围里,有的头只看附近几个像素,有的头关联较远区域,最后softmax出来的注意力分布自然就不一样。这些不同子空间的特征在拼接后进入后续卷积层,相当于给网络提供了多个视角的“输入素材”,特征表达的完备性会明显好于所有头都在同一个尺度上建模的方案。我用HMHA替换YOLOv12某个Neck节点里的特征融合方式后,检测头拿到的特征图通道确实比之前“干净”了,分类和回归分支的响应更集中。
2.2 子空间拆分如何拉开不同头的表达差异
子空间拆分是HMHA最核心的设计,但拆法有讲究。如果只是简单分组,每组仍然做全局注意力,那组与组之间还是容易趋同。要让子空间真正产生差异,每组必须有不同的“感知偏好”,常见做法是给不同组设置不同的窗口大小或不同的下采样比例。尺度越小的组,注意力范围越窄,适合捕捉小目标的局部纹理;尺度越大的组,注意力范围越广,适合建模大目标的全局上下文。我采用的配置是四组窗口分别为小、中、大、全局,这样从局部到全局都有对应的头去覆盖。
这个拆分方式可以类比成多镜头拍摄,广角镜头负责全景,长焦镜头负责局部特写,多个镜头合成一条丰富度极高的视频流。具体实现时,每个子空间内部仍然走标准的qkv投影和注意力计算,只是在注意力矩阵上叠加一个窗口mask,超出窗口范围的token一律置为负无穷,softmax之后权重自然归零。这样改动的计算成本很低,但表达差异非常明显。我从训练日志里统计过各头的注意力熵,HMHA的四个分组头熵值差距明显拉开,而普通MHSA的头熵值几乎挤在同一条水平线上。
2.3 通道重排为什么能降冗余
如果只做分组注意力,会引入一个副作用:组间信息隔离。每个子空间只在窗口内做特征的相互关联,不同子空间之间缺少信息交换通道,下采样或降维后这些组对应的通道会各说各话,相互之间的相关性反而可能失控。HMHA借鉴了轻量网络中通道混洗的思路,在分组注意力计算完之后,对输出特征在head维度上做一次重排,把相邻的通道打乱重组。
通道重排的本质是让不同子空间的信息在下一次计算前“碰个面”。比如第1组的第1个头原本输出位置在通道0和1,重排后变成和原第2组某个头相邻,后续的1×1投影层就能同时看到来自不同子空间的信息,从而学会跨子空间的组合模式。这个操作几乎不消耗额外参数和计算量,只是一次索引重排,但收益实打实。我在复现时做过对照,去掉通道重排后,mAP掉了0.4个百分点左右,说明这个看似不起眼的操作对最终精度有实质贡献。它和ShuffleNet的channel shuffle原理一致,只是HMHA将它应用在了head维度而不是卷积的group维度上。
3. 把HMHA接到YOLOv12 Neck的实操过程
3.1 先写一个可运行的HMHA模块
下面是一个简化但可以实际跑起来的PyTorch实现,包含了分组窗口mask和head维度上的通道重排。这个模块可以直接套进YOLOv12的Neck结构里,输入输出都是一个统一的B×N×C张量,N是特征图的H×W。为了代码清晰,我把一维窗口mask作为示例,实际使用时可以换成二维网格mask,逻辑完全相同。
import torch import torch.nn as nn import torch.nn.functional as F def make_window_mask(seq_len, window_size, device): """ 生成一个 (seq_len, seq_len) 的布尔mask, 只保留空间位置上距离不超过 window_size 的 token 对, window_size=0 代表不做窗口限制,即全局注意力。 """ if window_size == 0: return None idx = torch.arange(seq_len, device=device) dist = (idx[:, None] - idx[None, :]).abs() return dist <= window_size class HMHA(nn.Module): """ HMHA:分层多头注意力 把 num_heads 个注意力头分成 num_groups 组, 每组分别在自己的子空间窗口内计算注意力, 输出前对 head 维度做 channel shuffle,打破组间信息隔离。 """ def __init__(self, dim, num_heads=8, num_groups=4, window_sizes=(3, 7, 15, 0), qkv_bias=False, attn_drop=0.0, proj_drop=0.0): super().__init__() assert num_heads % num_groups == 0, "head数必须能被group数整除" assert len(window_sizes) == num_groups self.dim = dim self.num_heads = num_heads self.num_groups = num_groups self.heads_per_group = num_heads // num_groups self.head_dim = dim // num_heads self.window_sizes = window_sizes self.scale = self.head_dim ** -0.5 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) # 在 head 粒度上做 channel shuffle # 以 8 heads、4 groups 为例: # reshape => [[0,1],[2,3],[4,5],[6,7]] # 转置 => [[0,2,4,6],[1,3,5,7]] # flatten => [0,2,4,6,1,3,5,7] shuffle_index = torch.arange(num_heads).reshape( num_groups, self.heads_per_group ).transpose(0, 1).reshape(-1) self.register_buffer("shuffle_index", shuffle_index, persistent=False) def forward(self, x, seq_len): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim) qkv = qkv.permute(2, 0, 3, 1, 4) # (3, B, heads, N, head_dim) q, k, v = qkv.unbind(0) attn = (q @ k.transpose(-2, -1)) * self.scale # (B, heads, N, N) # 按group施加不同的窗口mask,实现子空间拆分 if self.window_sizes is not None: for g in range(self.num_groups): w = self.window_sizes[g] if w == 0: continue mask = make_window_mask(N, w, attn.device) h_start = g * self.heads_per_group h_end = h_start + self.heads_per_group attn[:, h_start:h_end] = attn[:, h_start:h_end].masked_fill( ~mask[None, None, :, :], float("-inf") ) attn = attn.softmax(dim=-1) attn = self.attn_drop(attn) x = (attn @ v) # (B, heads, N, head_dim) x = x.transpose(1, 2).reshape(B, N, self.num_heads, self.head_dim) # 通道重排:在head粒度上打乱顺序 x = x[:, :, self.shuffle_index, :] x = x.reshape(B, N, C) x = self.proj(x) x = self.proj_drop(x) return x提示:上面代码里我用一维距离近似了二维空间距离,目的是让mask逻辑更直观。真正接入2D特征图时,建议把序列索引还原成H×W网格坐标,再计算二维距离,这样窗口形状会更符合空间语义。
3.2 插入Neck的三种布局方案
模块写好了,接下来要决定放在Neck的哪个位置。我试过三种布局方案,各有优劣,下面直接给结论。
方案A:把HMHA插在PAN-FPN每次concat之后的那个融合卷积之前。这个位置特征是刚刚完成跨尺度拼接,通道冗余最严重的时候,HMHA在这里做一次注意力筛选和通道重排,能直接压掉一部分冗余信息,后面的融合卷积压力会小很多。优点是对整体结构改动小,缺点是每个尺度节点都要插,参数量和计算量累计起来比较大。
方案B:只在P3、P4、P5三个尺度的最终输出端各加一个HMHA。这种做法的思路是让注意力“把关”最终送给检测头的特征。优点是改动最收敛,速度影响小;缺点是Neck内部融合过程仍然有冗余,模型前半段的问题没有解决。我实际跑下来,方案B的精度提升明显弱于方案A。
方案C:只挑最大尺度和最小尺度两个节点插入HMHA,中等尺度保持原样。这个方案适合算力紧张的情况,因为P5分辨率低、计算量小,P3分辨率高、对小目标最敏感,两头都照顾到了。精度提升介于方案A和方案B之间,但训练速度损失最小。
三种布局的定位和算力开销对比如下:
| 布局方案 | 插入位置 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|---|
| A | concat后、融合卷积前 | 冗余处理最彻底,精度提升最高 | 参数量增幅较大,训练变慢 | 对精度要求高、算力宽裕 |
| B | P3/P4/P5最终输出端 | 改动最小,结构清晰 | 内部冗余未处理,提升有限 | 快速验证HMHA效果 |
| C | 仅P3和P5节点 | 平衡精度与速度 | 中等尺度特征未覆盖 | 有实时性要求 |
我个人建议先按方案C起步,跑通流程后在算力允许的情况下升级到方案A。这样既能快速看到HMHA带来的收益,也能留出调参的余地。
3.3 配置改动、参数量预算与训练策略
以一个中等规模的YOLOv12配置为例,Neck部分特征通道数通常在256左右,此时一个HMHA模块的参数量约等于qkv投影加输出投影,也就是约4倍通道数的平方,256通道时大约是0.26M参数。如果按方案A插入三个HMHA模块,总增加参数约0.8M,对于整体参数量在7M上下的模型来说,增幅不到12%,属于可以接受的范围。
在YOLO系列的yaml配置里,可以新增一个自定义模块名,然后像下面这样占位:
neck: - [-1, 1, "Conv", [256, 3, 2]] # concat操作之后插入HMHA - [-1, 1, "HMHA", [256, 8, 4]] # dim=256, heads=8, groups=4 - [-1, 1, "Conv", [256, 1, 1]]实际接入时要注意,YOLOv12的某些版本在Neck里用了C2f结构,我的做法是把HMHA作为C2f的替代或补充模块,放在concat之后、卷积之前,这样对原有特征流动路径的破坏最小。训练策略上我调整了几个关键点:一是学习率要适当调低,原来用0.01,加HMHA后我改成0.008,因为注意力模块对学习率更敏感;二是warmup轮数至少增加到原来的1.5倍,让qkv投影先稳定下来;三是drop path建议设置0.1左右,这个正则对防止注意力过拟合很有效。混合精度训练时HMHA的softmax和mask操作在FP16下是稳定的,但建议保持torch.autocast默认策略,不要把整个模块强行切到FP32。
4. 实验效果和踩坑实录
4.1 一组可以参照的实验对比
我先把一组参考数据放出来,它不是某个公开榜单上的绝对数字,而是我在一个内部检测场景下跑出来的相对走势。场景是工业零件表面缺陷检测,小目标和密集排列目标非常多,正好能放大HMHA在局部特征建模上的优势。训练统一用640输入,50个epoch,优化器SGD,损失函数和YOLOv12默认保持一致。
| 配置 | mAP50 | mAP50:95 | 参数量 | 单卡训练耗时变化 | 推理延迟变化 |
|---|---|---|---|---|---|
| 原始Neck | 52.3% | 34.8% | 7.1M | 基线 | 基线 |
| +SE模块 | 52.7% | 35.0% | 7.3M | +3% | +0.4% |
| +普通MHSA | 53.0% | 35.3% | 7.5M | +18% | +22% |
| +HMHA(8头4组方案C) | 53.4% | 35.9% | 7.4M | +6% | +4% |
| +HMHA(8头4组方案A) | 53.5% | 36.2% | 7.9M | +10% | +7% |
从数据上看,HMHA相比普通MHSA有两点明显优势:一是精度提升更多,二是计算开销更小。原因就是HMHA的窗口mask让大部分注意力计算局限在局部,不会像全局注意力那样产生大量无效的远距离交互。分类和回归分支的收敛曲线也显示,HMHA在训练后半段对验证集mAP的拉升更持续,说明子空间拆分确实增强了模型的泛化能力。
4.2 训练中常见的坑与排查方法
第一个坑是通道重排索引写错。我最初复现时直接把shuffle逻辑作用在特征通道维度上,而不是head维度上,结果训练时精度比基线还低了1.2个百分点。排查了两天才发现,原因是通道维度的随机重排会把同一个head内多个head_dim切片拆得七零八落,投影层根本学不到完整的子空间特征。正确的做法一定是在head粒度重排,每个head内部的head_dim保持连续,这点在代码里已经写清楚了。
第二个坑是显存占用不稳定。方案A在全分辨率特征图上做全局注意力那一组时,显存偶尔会跳高。我调整了window_sizes,把最大的窗口从0改为覆盖三分之二尺度的局部窗口,显存就稳定了。对于检测任务来说,完全全局的上下文在大目标上确实有帮助,但并非每个尺度节点都需要,P5层保留一个全局组就够了,P3和P4建议全部用局部窗口。
第三个坑是小数据集上不涨点。如果你的数据集只有几千张图,HMHA的分组结构反而会因为参数变多而容易过拟合。这时候可以试试先冻结Backbone,只训练Neck和检测头,或者把drop path调大到0.2。我自己在数据量偏小的另一个项目上遇到过类似问题,冻结Backbone后精度提升就恢复到合理范围了。
第四个坑是阈值类指标和回归指标的提升不同步。HMHA对分类置信度的分布影响比较明显,有时候mAP50涨了但mAP50:95没有同步涨,这是因为子空间拆分让分类分支学到的特征更锐利,但回归分支还没跟上。解决办法是把边框回归损失权重临时上调10%到20%,让回归分支在训练中后期重新匹配新的特征分布。
4.3 参数调节的顺序与建议
调节HMHA参数时,建议按下面的优先级顺序来,不要一上来就同时改好几个维度。
先定分组数。分组数决定了子空间拆分的颗粒度,我试过2组、3组、4组、8组,4组在YOLOv12的Neck场景下综合表现最好。2组时子空间差异不够大,局部和全局的区分度不高;8组时每组只有一个头,头的感受野切换太陡,训练稳定性反而下降。
再定每个组的窗口配置。这个要以数据集中目标的尺度分布为依据。如果你的检测目标普遍较小,小窗口那组可以多分配一个组;如果目标尺寸跨度大,就把中间窗口的覆盖范围拉开。窗口配置对效果的影响比分组数更直接,值得多试几轮。
最后调drop path和注意力dropout。这两个参数决定注意力权重的平滑程度,数值过大容易欠拟合,过小则容易在训练后期震荡。我这边drop path取0.1、attn_drop取0.0时效果最好,但不同数据集的最优值会有浮动,建议用小幅网格搜索。
5. 还能往哪个方向扩展
5.1 轻量化和剪枝
HMHA的一大好处是结构规整,分组和头数的设定天然适合做结构化剪枝。训练完成后,可以统计每个注意力头的有效注意力覆盖范围和熵值,把那些熵值持续偏小、对精度贡献低的头剪掉,再把shuffle_index重新生成一次。剪掉的头对应的通道会整体移除,从而直接减少后续投影层的输入维度。这样可以在几乎不掉点的情况下压缩模型,尤其适合要部署到边缘设备的场景。
5.2 与检测头联动的改法
除了在Neck里使用HMHA,还可以把HMHA输出的多组注意力权重直接引向检测头的分类分支。具体做法是把最后一层HMHA每个group的输出单独保留,分别经过轻量卷积后和检测头的分类特征做相加。这样分类分支就能同时借助局部窗口的小目标信息和全局窗口的上下文信息,对小目标分类有明显帮助。这个改法会增加少量计算,但效果可以和Neck内的HMHA叠加,我试过在密集场景下小类别的recall提升了约1.5个点。
5.3 迁移到其他多尺度任务
HMHA并不绑定YOLOv12,任何需要多尺度特征交互的任务都可以用同样的思路。图像分割中,可以将它插入FPN结构的顶层输出,让不同感受野的分组头分别捕捉器官边界和全局结构;多目标跟踪里,可以在ReID分支前插入小窗口的HMHA,强化局部表观特征。关键点依然是子空间设计和通道重排这对组合,前者保证多样性,后者保证信息流通。换到其他任务时,只需要重新设计窗口大小配置,主结构不用动。
最后说一个我在复现过程中感触最深的细节。第一次把HMHA接进Neck时,我以为重点全在子空间拆分上,注意力全放在窗口mask的逻辑上,结果忽略了channel shuffle这一步。等到对照实验出来,去掉shuffle的版本比完整版掉了0.4个点,我才意识到道理其实很简单:把一个团队分成几个小组还不够,关键是小组之间要有信息交换的机制,否则各干各的,最后合到一起还是两张皮。这个道理放在模型结构里,就是那步看似不起眼的通道重排。做模型改进不怕想法复杂,怕的是理所当然地认为“小细节不影响大局”。HMHA这个改进能在我这个场景里稳定涨点,靠的正是对这两个看似基础的问题较真。