102、C2PSA 替换为 PSAMask 空间注意力掩码的代码实现与训练稳定性分析

102、C2PSA 替换为 PSAMask 空间注意力掩码的代码实现与训练稳定性分析

一、从一次诡异的loss震荡说起

上个月调YOLOv11n在VisDrone上的检测任务，C2PSA模块在训练到第80个epoch时突然loss开始剧烈抖动，val mAP从0.42掉到0.38又弹回0.41，反复横跳。当时我盯着tensorboard看了半小时，第一反应是学习率策略出了问题，但检查cosine annealing曲线完全正常。后来逐层打印梯度范数，发现C2PSA内部的self-attention分支在深层的梯度方差异常大——某些batch的梯度范数能达到其他batch的10倍以上。

这个现象让我想起之前读的一篇论文《PSAMask: Learning Spatial Attention Masks for Stable Training》，里面提到用可学习的空间注意力掩码替代自注意力机制，能显著抑制梯度震荡。于是决定把C2PSA整个替换成PSAMask模块，顺便验证一下在YOLOv11这种anchor-free检测器上的效果。

二、C2PSA的问题到底在哪

先看C2PSA的核心结构（这里不贴图，直接说代码逻辑）：它把输入特征分成两路，一路走常规卷积，另一路走多头自注意力（MHSA）。问题出在MHSA的softmax注意力权重上——在训练初期，特征图的空间位置之间相关性很弱，softmax输出的分布接近均匀分布，梯度回传时每个位置的梯度几乎相等，导致参数更新方向模糊。随着训练进行，某些位置逐渐获得高注意力权重，但这些权重的梯度方差会突