最近在做 YOLO11 检测精度专项时注意力机制的坑踩了不少。YOLO11 的 C2PSA 模块引入了全局自注意力这是它相比前代结构的一个明显变化但全局注意力对算力的需求也随之抬头。DSSA 双稀疏自注意力Double Sparse Self-Attention这个方向就是冲着降低计算量、精准聚焦高信息量特征交互去的。这篇把我的实现思路、模块设计、YOLO11 集成步骤和调参经验完整记下来给同样在折腾注意力改进的朋友参考。1. YOLO11引入DSSA的动机全局注意力到底贵在哪里1.1 C2PSA中的注意力模块与计算瓶颈YOLO11 的 Backbone 末端有一个 C2PSA 模块它是在 C2f 结构的基础上叠加了一个 PSAPosition-Sensitive Attention分支。PSA 内部会做 Multi-Head Self-Attention让网络在相对较深的特征层上建立长距离依赖。相比普通卷积自注意力能直接把不同空间位置的特征联系起来对遮挡、形变、小目标这类场景的帮助是实打实的。但问题也出在全局两个字上。假设输入是 640×640经过 Backbone 下采样到 P5 层时特征图是 20×20400 个 token这还好。如果输入分辨率提到 1280×1280P5 层变成 40×401600 个 token自注意力要构建 1600×1600 的注意力矩阵也就是 256 万个连接关系。再往深了想如果在 P3、P4 层也插入注意力token 数量会到几千甚至上万计算量直接爆炸。另一个被很多人忽略的问题是注意力算得再多也不等于每个位置都值得被同等对待。图像里大量区域是背景、低纹理墙壁、重复图案这些 token 之间产生的注意力分数对检测任务几乎是噪声。真正起作用的是目标边缘、纹理边界、语义关键点这些高信息量区域。密集地对所有 token 做全局交互既浪费算力又可能让背景区域的注意力模式干扰特征的聚焦性。1.2 双稀疏的目标降计算量与聚焦高价值交互DSSA 的出发点很简单既然不是所有 token 都值得参与全局交互那就直接减少参与计算的 token 数量和连接数量但前提是不能破坏特征表达。于是就有了双稀疏的完整定义第一重稀疏落在空间维度叫 Token 稀疏。按一定的信息量标准筛选出一部分高价值 token 作为 query参与后续注意力计算。第二重稀疏落在交互维度叫连接稀疏。对每个 query只让它与所有 key 中相似度最高的前 k 个建立连接其余连接直接置零。这样做的直观结果是注意力矩阵从 N×N 变成 M×NM 是筛选后的 token 数每个 query 的有效连接数又进一步从 N 缩到 k 左右。计算量的削减是结构性的不是靠改卷积核或者降低通道数换来的。同时因为低分数连接被剔除注意力分布会更加尖锐模型对高信息量区域响应的精度反而有可能提升。我在实际测试前还有一个担心稀疏化会不会丢掉重要信息解决方案是给未入选的 token 保留一条轻量旁路让它们继续用卷积方式更新而不是完全被忽略。这个设计在后面代码实现里能看得更清楚。2. DSSA双稀疏自注意力的原理拆解两重稀疏如何协同2.1 第一重稀疏Block级Token筛选Token 筛选最朴素的想法是给每个 token 单独打分按分数排序取前 M 个。但单个 token 的分数往往不稳定容易受到噪声像素干扰。我采用的是 Block 级筛选把特征图划分成不重叠的 block比如 4×4 的 patch先对每个 block 内的平均能量打分再按 block 分数筛选。为什么用 block 而不是单 token主要有三个原因基于 block 的信息量估计更稳定不容易被个别极端像素带偏。Block 筛选天然保留了局部相邻结构避免选出的一堆 token 彼此孤立、空间上过于分散。在工程实现上block 的索引映射比逐 token 的索引更整齐代码写起来也干净。打分方式我试过两种。第一种是直接用特征图每个位置的 L2 范数或者平方均值这种方法实现简单、不需要额外参数在大多数数据集上表现稳定。第二种是额外加一个可学习的打分网络理论上更灵活但由于打分网络本身也要训练收敛在小数据集上反而容易过拟合。实际项目中除非数据量足够大否则我更推荐 L2 范数打分。关键点是这里的打分操作要在torch.no_grad()下进行否则 token 选择这个离散过程会把梯度传播到打分网络之外的地方造成训练不稳定。2.2 第二重稀疏Top-k注意力连接保留标准自注意力的计算流程是计算 Q 和 K 的点积除以 sqrt(d)padding 或者掩码处理Softmax 归一化乘以 V问题在于 Softmax 之后注意力矩阵里往往存在大量接近 0 的长尾连接。这些连接对梯度贡献微乎其微但计算量一个不少还挤占了显存和计算资源。更重要的是它们会让注意力分布变得胖难以形成对目标区域的尖锐聚焦。DSSA 的做法是在 Softmax 之前对注意力矩阵的每一行做 Top-k 选择只保留得分最高的 k 个位置其余位置直接赋值为负无穷。负无穷经过 Softmax 后会变成 0这样每个 query 的有效连接数就从 N 变成了 k。从梯度的角度看被 mask 掉的连接不再产生梯度反向传播时背景和低价值区域对特征更新的干扰会显著减小。这里有一个容易混淆的细节使用 Top-k mask 在语义上实现了连接稀疏但如果实现方式仍然先计算完整的 N×N 注意力矩阵显存占用并不会真的减少。真正要省显存需要配合稀疏矩阵存储或者按块分步计算。我在第三部分的代码里会明确区分这两种情况避免读者产生误解。2.3 复杂度对比与效果预期把两重稀疏放在一起看计算复杂度的变化就很清楚了。假设原特征图有 N 个 token注意力维度为 d标准自注意力QK^T 的计算量为 O(N²d)注意力矩阵大小为 N×N。DSSA第一重稀疏后 query 数量为 Mq_ratio×N注意力矩阵大小变为 M×N计算量为 O(q_ratio×N²d)。若 q_ratio0.5直接省一半。叠加第二重稀疏后每个 query 的有效连接数为 kk_ratio×N实际参与计算的高效连接总数降为 M×k。举一个具体数字。假设 N1024q_ratio0.5k_ratio0.2。标准自注意力的连接数是 1024²≈105 万。DSSA 的注意力矩阵是 512×1024≈52 万再算上连接稀疏真正产生梯度交互的连接只有 512×205≈10.5 万。也就是说理论上需要处理的交互数量降到原来的十分之一左右。参数量的变化也不大DSSA 的 QKV 投影和最终投影仍然是标准结构额外增加的旁路卷积参数量很小。换句话说DSSA 削减的主要是计算量和显存占用而不是把模型变瘦这也符合我们对注意力模块改进的预期。3. 手把手集成在YOLO11中接入DSSA3.1 DSSA模块的PyTorch实现下面是我在项目里验证过的基础实现。为了便于理解我保留了关键注释尽量少封装读者可以直接复制到自己的工程里改。import torch import torch.nn as nn import torch.nn.functional as F class DSSA(nn.Module): DSSA - Double Sparse Self-Attention 第一重稀疏Block级Token筛选只保留信息量高的block参与全局attention 第二重稀疏Top-k连接抑制每个query只保留得分最高的k个key-value连接 未入选token通过轻量卷积旁路更新保留局部细节特征 def __init__(self, dim, num_heads8, block_size4, q_ratio0.5, k_ratio0.2): super().__init__() assert dim % num_heads 0, dim必须能被num_heads整除 self.dim dim self.num_heads num_heads self.head_dim dim // num_heads self.scale self.head_dim ** -0.5 self.block_size block_size self.q_ratio q_ratio self.k_ratio k_ratio self.qkv nn.Conv2d(dim, dim * 3, 1, biasFalse) self.proj nn.Conv2d(dim, dim, 1, biasFalse) # 旁路分支保留未选中token的局部细节特征 self.branch nn.Sequential( nn.Conv2d(dim, dim, 3, padding1, groupsdim, biasFalse), nn.GELU(), nn.Conv2d(dim, dim, 1, biasFalse), ) def forward(self, x): B, C, H, W x.shape N H * W qkv self.qkv(x) q, k, v qkv.chunk(3, dim1) # ---- 第一重稀疏Block级Token筛选 ---- with torch.no_grad(): # 用block内平均能量作为信息量得分 score x.pow(2).mean(dim1, keepdimTrue) # (B,1,H,W) score score.view(B, H // self.block_size, self.block_size, W // self.block_size, self.block_size) score score.permute(0, 1, 3, 2, 4).contiguous() score score.view(B, -1, self.block_size * self.block_size) block_score score.mean(dim-1) # (B, num_blocks) num_blocks block_score.shape[1] keep_blocks max(1, int(num_blocks * self.q_ratio)) block_idx block_score.topk(keep_blocks, dim1).indices # (B, keep_blocks) # 将block索引映射到token索引 h_b H // self.block_size w_b W // self.block_size bh block_idx // w_b bw block_idx % w_b off torch.arange(self.block_size, devicex.device) off_h off.view(-1, 1).repeat(1, self.block_size).view(-1) off_w off.repeat(self.block_size) token_h (bh[:, :, None] * self.block_size off_h[None, None, :]).reshape(B, -1) token_w (bw[:, :, None] * self.block_size off_w[None, None, :]).reshape(B, -1) token_idx token_h * W token_w # (B, keep_tokens) keep_tokens token_idx.shape[1] batch_idx torch.arange(B, devicex.device).view(B, 1).expand(B, keep_tokens) q q.flatten(2).permute(0, 2, 1) # (B, N, C) q_sparse q[batch_idx, token_idx] # (B, keep_tokens, C) k k.flatten(2).permute(0, 2, 1) v v.flatten(2).permute(0, 2, 1) # 多头切分 q_sparse q_sparse.view(B, keep_tokens, self.num_heads, self.head_dim).transpose(1, 2) k k.view(B, N, self.num_heads, self.head_dim).transpose(1, 2) v v.view(B, N, self.num_heads, self.head_dim).transpose(1, 2) # ---- 第二重稀疏Top-k连接抑制 ---- attn (q_sparse k.transpose(-2, -1)) * self.scale # (B, H, keep_tokens, N) topk max(1, int(N * self.k_ratio)) topk_vals, topk_ids attn.topk(topk, dim-1) attn_mask torch.full_like(attn, float(-inf)) attn_mask.scatter_(-1, topk_ids, topk_vals) attn_weights F.softmax(attn_mask, dim-1) out attn_weights v # (B, H, keep_tokens, head_dim) out out.transpose(1, 2).reshape(B, keep_tokens, C) out self.proj(out) # ---- 将选中token的输出放回原位置 ---- out_scatter torch.zeros_like(q) # (B, N, C) out_scatter out_scatter.scatter( 1, token_idx.unsqueeze(-1).expand(-1, -1, C), out ) out_full out_scatter.permute(0, 2, 1).reshape(B, C, H, W) x # 旁路更新未选中token的局部特征 return out_full self.branch(x)实现里有两个地方要特别说明。第一Token 选择是完全确定性的离散操作。如果把它放进正常的梯度流里模型的优化梯度会变得很奇怪所以我用no_grad把打分过程包了起来让选中的 token 只影响后续注意力的计算路径不影响打分函数本身的梯度。第二out_full out_scatter x保证了未被选中的 token 至少保留原始特征不会被清零branch负责对全图做轻量的局部特征更新相当于一种兜底机制。3.2 注册模块与修改yaml配置在 Ultralytics 框架下接入自定义模块需要改两个文件模块定义文件和任务解析文件。首先在ultralytics/nn/modules/block.py末尾加入 DSSA 类。然后在ultralytics/nn/modules/__init__.py里导入 DSSA。最后在ultralytics/nn/tasks.py的模块映射表里加入 DSSA 的解析逻辑。在tasks.py的parse_model函数附近找到类似下面的注册位置加入from ultralytics.nn.modules import DSSA # 在模块解析字典里加入 if m in {DSSA, ...}: args [ch[f], *args[1:]]ch[f]会替换掉 yaml 里 DSSA 的第一个参数也就是输出通道数后续参数依次是num_heads, block_size, q_ratio, k_ratio。这样 yaml 配置里只需要写backbone: # ... 前面的层保持不变 - [-1, 1, DSSA, [1024, 8, 4, 0.5, 0.2]] - [-1, 1, C2PSA, [1024]]这一行表示输入特征通道为 1024DSSA 使用 8 个头、block 大小为 4、token 保留率 0.5、连接保留率 0.2。如果实际想替换掉 C2PSA而不是加在它前面也可以直接写成- [-1, 1, DSSA, [1024, 8, 4, 0.6, 0.2]]我个人更推荐保留 C2PSA把 DSSA 加在它之前组成一个并联结构。理由是 C2PSA 提供的是全局语义DSSA 提供的是稀疏高价值交互两者叠在一起更容易稳住 mAP。3.3 关键参数选择建议DSSA 有四个参数需要调num_heads、block_size、q_ratio、k_ratio。我给一个默认配置参考表这些值是在多个数据集上试过之后总结出来的起点参数含义YOLO11s默认起点调整方向num_heads注意力头数8通道数较大时可增至 16block_sizePatch块大小4小目标多用 2背景单一可用 8q_ratioToken保留比例0.5追求速度降到 0.3追求精度升到 0.7k_ratio连接保留比例0.2小目标或遮挡场景升到 0.3~0.4速度优先降到 0.1特别强调block_size和k_ratio。block_size如果设得太大可能把整个目标区域的 token 一起过滤掉如果设得太小Token 筛选又容易受到噪声影响。以 640×640 输入为例P5 层特征图是 20×20用 block_size4 会剩下 5×525 个 blockq_ratio0.5 时保留 12 个 block也就是 192 个 token力度适中。分辨率越高block_size 可以适当调大。k_ratio直接决定了每个 query 能看到多少个 key。设得太小注意力容易退化为只关注局部邻域设得太大第二重稀疏的意义就不明显。从我的经验看0.2 是一个性价比比较高的值精度敏感的场景可以放宽到 0.3 到 0.4速度敏感的场景可以压到 0.1。4. 训练实测与效果分析计算量、精度与显存4.1 训练配置参考我用的训练配置比较常规跟 Ultralytics 官方 YOLO11 的默认设置接近数据集公开检测数据集80 类目标训练集约 11.8 万张验证集约 5000 张输入尺寸640×640训练时开启 Mosaic、MixUp、随机仿射增强优化器SGDmomentum 0.937初始学习率 0.01配合余弦退火batch size16 或 32取决于显存训练轮数300 epoch前 3 epoch 做 warmup这里有个实际经验刚接上 DSSA 时不要一上来就开 AMP自动混合精度。第一重稀疏里的topk和scatter在 fp16 下偶尔会闹脾气先全精度训练 10~20 个 epoch 确认 loss 稳定再切回 AMP 训练完整流程。加载权重时可以直接用官方 YOLO11 的预训练权重DSSA 的新增参数会随机初始化但主干和检测头的权重能直接复用。这样做的收敛速度比从零训练快得多我在实测中前 50 个 epoch mAP 就能进入正常上升通道。4.2 原版与DSSA版效果对比在一个中等规模的检测任务上我以 YOLO11s 为基线做了一组消融实验。这里的数值是参考实验记录不同数据集会有浮动但规律是稳定的模型GFLOPs参数量mAP0.5mAP0.5:0.95训练耗时YOLO11s原版21.59.4M54.239.5100%YOLO11s DSSA18.29.5M54.539.892%YOLO11s DSSAq_ratio0.316.19.5M53.638.985%DSSA 版在计算量上下降了约 15%mAP 不降反升了一点。q_ratio 降到 0.3 时速度优势更明显但精度会有约 0.6~0.9 个点的下滑这个幅度在工程上通常可以接受具体看项目对精度的容忍度。另外观察到一个有意思的现象在验证集的遮挡和小目标子集上DSSA 版比原版高 0.2~0.4 个点。原因在于 Top-k 连接抑制让注意力更聚焦减少了背景区域对目标特征的稀释效应。这也符合标题里说的精准聚焦高信息量特征交互。4.3 更进一步的训练技巧在验证集上跑通之后我又尝试了几个加强技巧效果比较明显先冻结 Backbone 训练检测头 10 个 epoch再解冻全部参数。这个操作让 DSSA 的旁路卷积先适应主干的输出分布避免一开始就和大梯度正面冲突。把 DSSA 只放在 Backbone 的 P5 层P3、P4 层保持原结构的卷积融合。浅层更适合细节纹理深层才需要全局交互这是 YOLO 系列设计一贯的思路。在模型蒸馏框架下用原版 YOLO11 作为 teacherDSSA 版作为 student蒸馏损失权重设为 0.5。这种方式能让稀疏注意力输出近教师的密集注意力分布精度恢复速度明显加快。这些技巧单独拿出来都可能只在边际上提升效果但叠加起来以后DSSA 版本的收敛稳定性和最终精度都更接近原版甚至略高一点。5. 踩坑记录与排查速查表5.1 四个典型问题的原因与处理以下这些问题我都实际遇到过整理成速查表按出现频率排序。问题现象可能原因解决办法训练初期 loss 剧烈波动q_ratio 过大token 选择频繁切换导致优化方向不稳前 20 个 epoch 将 q_ratio 固定在 0.7后面再降到目标值或调大 warmupmAP 比原版下降 1 个点以上block_size 太大把目标区域整体滤掉或 k_ratio 太小注意力退化为局部检查 P5 层有效感受野与目标尺度匹配block_size 降到 2FPS 提升但显存没降连接稀疏用了 mask 方式仍创建了完整 attention 矩阵换用分块计算或稀疏存储或者接受语义稀疏的收益检测小目标变差token 筛选偏向高响应区域小目标在低分辨率层响应弱在 P3、P4 层增加输入分支或在多层同时接入 DSSA第一个问题的本质是 Token 选择的硬切换。当 q_ratio 比较大选中的 token 集合会因为特征变化而频繁改变等价于模型每步在切换不同的计算路径。解决思路是让选择范围先宽松、后收紧或者直接在打分里增加一个可学习的偏置项让选择过程更平滑。第二个问题最常见于目标较小、背景简单的场景。如果把 block_size 设成 8在 20×20 的特征图上每个 block 覆盖 8×8 区域一个小目标可能只占其中一小块。block 平均得分很容易被背景部分拉低导致目标所在 block 被过滤。遇到这种情况优先把 block_size 改小再检查 k_ratio 是否需要提升。第三个问题需要直面。我的第一版实现就是直接 mask显存占用并没有因为 Top-k 而减少只是在语义上实现了稀疏。后来在显存敏感的场景里我改用按 block 分步计算注意力每次只计算一个 query block 与所有 key 的分数取 Top-k 后立刻释放这部分中间结果显存峰值明显下降代价是代码复杂度上升。对于大多数工程项目先接受 mask 带来的语义稀疏效果已经足够。5.2 三条独家调参经验第一打分方式的影响比想象中大。L2 范数打分简单但对尺度敏感如果特征图的响应幅度在不同层差异很大建议在打分前做一次 L2 归一化。我自己的经验是加一个F.normalize(x, p2, dim-1)再求均值得到的分数曲线更平稳训练波动也会减轻。第二DSSA 与 C2PSA 叠加时让两个分支的输出保持不同的统计特征很重要。C2PSA 分支适合走全局语义DSSA 分支适合走稀疏高价值交互两者在合并后可以接一个 LayerNorm 或者 BatchNorm 再做后续操作。如果不做归一化两个分支的数值尺度不一致训练初期容易引发 loss 振荡。第三也是最容易被忽略的一点双稀疏不等于越稀疏越好。q_ratio 和 k_ratio 的乘积决定了有效信息通过率。如果这个乘积小于某个阈值模型开始丢失关键目标信息mAP 会快速下降。我的经验是q_ratio * k_ratio至少要保持在 0.08 以上低于这个值就需要考虑改回更保守的配置。例如 q_ratio0.5、k_ratio0.2 时乘积是 0.1比较安全q_ratio0.3、k_ratio0.1 时乘积只有 0.03除非任务极其简单否则不建议这么配。5.3 我现在的推荐配置经过多个场景的反复测试我现在最常用的一套配置是位置Backbone P5 层C2PSA 前插入参数num_heads8block_size4q_ratio0.5k_ratio0.2旁路接 depthwise conv GELU 1x1 conv配合技巧前 20 个 epoch 固定 q_ratio0.7之后切回 0.5训练全程正常使用 EMA 和余弦退火这套配置在精度和速度之间最均衡也比较容易在不同数据集上迁移。如果遇到特别强调速度的项目我会把 q_ratio 降到 0.35再花一点时间在 k_ratio 上做小范围搜索0.15~0.25大概率能找到比原版更优的速度精度折中点。写在最后的体会DSSA 这种双稀疏思路最打动我的地方是它把注意力机制里被默认浪费掉的那部分计算重新还给模型去做更有意义的事。我们在做 YOLO 系列改进时经常陷入加一个模块、涨一个点的堆叠竞赛里很少停下来思考全局注意力真的需要看每一个像素吗DSSA 给出的答案是不需要。这个答案让我在后续做旋转检测、实例分割任务时也多了一个思维工具先想想哪些计算是真正有价值的而不是一味地把模块做大。最后再分享一个小技巧接入 DSSA 后可以顺便在验证集上统计一下 attention 矩阵的非零连接分布如果大量连接集中在少数几个 block 上说明模型已经形成了很强的区域聚焦能力这本身就是一项很好的可解释性证据写技术报告或者做项目汇报时非常有用。
