注意力机制的演进已经过了堆砌结构的阶段真正有价值的改动都藏在“最朴素的直觉”里。Kimi 团队这篇关于 Attention Residuals 的论文标题看起来平淡但把“残差”和“注意力”这两个关键词放在一起懂行的人马上会精神起来。它想解决的是注意力机制里一个很少有人正面处理的问题当注意力分布逐渐走向极端模型的信息提取能力其实是被自己削弱掉的。这篇解读我会按“问题动机 → 设计空间 → 工程落地 → 踩坑排查”的顺序拆开讲适合正在做注意力变体实验、或者想在自己模型上引入稳定收益结构的研究人员和工程师参考。1. 注意力机制为什么需要“残差”先从“坍缩”说起1.1 标准注意力不是你想的那样稳定Transformer 里的自注意力数学上就三行Q 和 K 做点积除以根号 dsoftmax 成权重再加权求和 V。这个机制在绝大多数场景下表现很好但它有一个被默认接受、其实很危险的特性softmax 会把输入向量的差异放大成概率差异一旦某些 token 的注意力分数比其他 token 高出一点经过指数运算后权重会瞬间向这些 token 集中。我见过太多人调模型时碰到的诡异现象深层网络训练到一半注意力图变成了几条又细又亮的线几乎所有 token 只看跟自己最近的两三个位置或者死盯着某个语义中心词不放。这不是偶然而是注意力坍缩的典型症状。序列越长越明显因为点积分数的方差会随着参与计算的元素增多而不稳定softmax 的输出越来越接近 one-hot 分布。这种坍缩带来的后果很直接模型名义上有全局建模能力实际只会用一小截上下文长文本里的“大海捞针”任务表现一言难尽逻辑推理时也容易漏掉远处的约束条件。过去很多解法都在打补丁比如温度缩放、flash attention 里的在线归一化优化、或者给注意力分数加上相对位置偏置如 ALiBi。这些手段缓解了分布极端化但都没有把“注意力本身退化”当成一个独立的建模问题来对待。1.2 残差思想从 ResNet 到注意力模块的迁移逻辑深度学习里遇到“退化”问题最先想到的应该是残差结构。ResNet 时代的核心洞见是与其让网络硬学一个复杂的映射 H(x)不如让它学 F(x) H(x) − x再在输出端加回 x。这样一来网络不需要从头记住恒等映射只要在已有表示上做增量调制就够了。这个思路让深层网络的梯度、表达能力、训练稳定性同时受益。注意力机制其实也有对应的“退化”问题。标准注意力的输出是对 V 的加权求和如果注意力权重坍缩成 one-hot那输出就直接退化成某个位置本身的向量multi-head 的多样性也跟着没了。换句话说注意力模块的“恒等映射”应该是保留每个 token 自身的原始信息就像 ResNet 里输出端必须有一条 x 的旁路。很遗憾标准注意力里恰恰没有这条旁路每个 token 所能获取的信息完全由归一化过的权重说了算。所以 Attention Residuals 这个概念本质上是把 ResNet 的问题意识搬进了注意力计算内部在注意力得分、注意力权重、或者注意力输出这三个关键位置上显式地保留一条来自“自身”的恒等路径让注意力不会因为归一化而弄丢原始的 token 身份信息。这个迁移逻辑听起来简单但选择在哪一个位置插入残差、用什么方式保底、如何和 mask 与归一化兼容就是论文真正在做文章的地方。2. Attention Residuals 的设计空间三种可能的落地形态2.1 形态一对注意力分数logits施加结构先验第一种实现方式是在 softmax 之前的 logits 上直接加一个残差矩阵 R让注意力分数变成 QKᵀ/√d R。这个 R 可以是固定的也可以是可学习的最常见的选择是把单位阵 I 的缩放版本作为先验。这个设计想表达的逻辑是什么本质上它是在给“自己 attending 到自己”这件事一个固定的偏置加分。我把它理解成给注意力机制安装了一个保安不管你最终对全局做了多少加权组合至少你要保证不会彻底忽略掉当前位置的信息。这对深层网络特别有意义因为多层叠加之后token 的原始语义经过反复变换已经非常抽象如果没有自连接保底浅层学到的部分特征可能在深层被注意力机制直接清洗掉。这类残差与已有工作的区别值得注意。位置编码偏置比如 ALiBi是在 logits 上根据两个 token 的相对距离加一个负向偏置它的目的是约束“远距离关系”而 Attention Residuals 里加的 I 矩阵修正的是“同一位置的恒等路径”两者不冲突甚至可以叠加使用。实现层面也干净R 通常是稀疏的几乎不增加显存访问。2.2 形态二对注意力权重做凸组合给行归一化留一条恒等路第二种形态更直接在 softmax 之后把注意力权重矩阵 A 和某个残差矩阵 R 做凸组合得到新的权重 A。如果 R 取单位阵 I就是给每个位置固定分配一部分注意力给自己用纯数学语言写就是A (1 − α) · softmax(QKᵀ / √d) α · I这个公式的好处是直观α 相当于一个“自我注意力比例”控制模型保留多少基础的自身路径。你觉得 α 是 0.05 还是 0.2取决于任务里自身信息的重要程度。如果任务本身几乎完全依赖长程检索α 太大反而会拖后腿。这里要坚持一个原则加完残差之后注意力权重必须依然满足行归一化。凸组合天然满足但如果你是直接在 softmax 的结果上加上一个未缩放的单位阵那这一行的概率总和会变成 1 α后面接的因果采样和 KV Cache 处理都会出问题。常见做法是利用凸组合保证归一化依旧成立代码里一行就能实现却常常被人忽略。2.3 形态三在注意力输出上做门控残差第三种形态把残差的位置挪到了输出端即注意力机制的输出不是一个纯 V 的加权和而是原始输入 x 和注意力输出的门控混合。这一步和 Transformer 原本的 Add Norm 有相似之处但角度不同原来的 Add 是在子层外做残差连接而这里是在注意力计算内部就把原始输入作为一份成分混进去。这种设计的核心动机是增强模块的“保底能力”。深层 Transformer 里靠 LayerNorm 和残差连接外加一个注意力子层理论上信息是可以跨层流动的但实际上子层内部多次线性变换后梯度已经变得非常碎片化。如果能在子层内部再加一道旁路相当于给梯度增加了一条更短的传播通道。2.4 多头分工与层间残差链再往上走一个层次Attention Residuals 还可以从“头”和“层”两个维度展开。多头注意力原本的设计意图是让不同的头捕获不同的关系模式但坍缩发生时多个头会趋同于同一个极端分布分工就形同虚设。一种直接的扩展是给每个头设置独立的 α让一部分头偏向自我保底另一部分头专注全局检索。这种“分工式残差”引导不同头走向不同的功能角色而不是放任它们自由趋同。在层间维度早期 Transformer 层输出的注意力信息经过后续层后容易被更高阶的特征覆盖。论文标题里的“Residuals”如果放在层间视角下还能理解成跨层的注意力跳跃连接——类似把浅层的注意力分布直接送到深层去参与计算。这类做法的风险在于实现复杂度上升对推理框架的改造比较大所以通常是三种基础形态外的进阶方向。3. 从论文构思到工程实现关键细节与实操参数3.1 先线性插值还是先做归一化两种公式的等价性与差异这里有个非常容易踩坑的细节把残差加在 softmax 之前和加在 softmax 之后数学上并不总是等价。加在 softmax 之前的那条路残差矩阵会和 QKᵀ 一起参加指数运算于是残差项的影响其实是整个分数分布的函数带着复杂的非线性耦合加在 softmax 之后做凸组合则是在已经形成的概率分布上做线性混合行为直白可预期。我在实践中更推荐先想清楚“你想让残差影响模型的方式是哪种”。如果你希望残差能动态对抗 logits 方差过大导致的坍缩那放在 softmax 之前更直接如果你只想要一个稳定、易解释、可退化为标准注意力的结构那放在 softmax 之后做凸组合更省心。论文如果只讲一种实现你理解之后一定要自己把这个对照组跑一遍因为两种方式在深层网络里的表现差异往往比想象中更大。推荐把两种公式明确写成# 形态一:logits 残差 logits (q k.T) / math.sqrt(d) residual_bias prob torch.softmax(logits, dim-1) # 形态二:权重凸组合 prob (1 - alpha) * torch.softmax((q k.T) / math.sqrt(d), dim-1) alpha * identity_matrix3.2 超参数怎么设α 的初始化、退火与自适应对凸组合形态来说α 是最关键的超参数。我见到不少复现党上来就拍脑袋设 α 0.1训完说提升太小这往往是因为 α 设置不合适。按我的经验α 的初始值要跟训练阶段绑在一起。模型刚起步时注意力分布没有成型α 可以稍大一些相当于给模型一个比较强的“自我保底”先验帮助它稳定训练训练中后期各类模式逐渐成熟如果 α 还维持很大模型会变得“短视”不善于检索远程信息。一个常见做法是让 α 从 0.1 或 0.15 开始线性退火到 0.02 左右或者干脆通过一个可学习的标量来自适应。可学习的好处是省去调参但要注意对可学习标量做正则否则它容易滑到某个极端值而不是停留在最优区间。我通常还会配合观察每一层的 attention entropy。如果某几层的 entropy 明显比其他层低说明这些层已经进入坍缩状态应该给它们更大的 α这就是“按层自适应”的朴素实现。3.3 兼容性问题FlashAttention、因果掩码与 KV Cache凡是改动注意力计算的方案我都建议先考虑它跟 FlashAttention 和因果掩码兼容吗Attention Residuals 在 logits 上做文章的话天然兼容在线 softmax因为只是在分数矩阵上加一个偏置但如果在权重矩阵层面做凸组合FlashAttention 的分块计算会遇到麻烦因为我们要把分块 softmax 的结果直接和单位阵混合这破坏了它标准的 online softmax 流程需要额外遍历一次或单独处理对角块。因果掩码的影响也不容忽视。做语言模型时必须保留 causal mask那么残差单位阵 I 在对角线以下的部分其实是无效的因为那些位置被 mask 掉了。正确的姿势是让残差矩阵只在对角线及其合法范围内存在否则等于从句把未来位置的数值泄漏进来了。这个 bug 很隐蔽报错不会出现但训练效果会莫名变差。KV Cache 层面小 α 下的残差几乎不引入额外的内存开销但如果做多头独立 α缓存里就要多存一份每头的 α 向量。虽然量很小但在推理框架里这个额外的元数据会带来一点调度开销尤其对那种一次推理成千上万个请求的线上服务。工程上最好把它压进浮点常量或者量化存储别在每步解码时都动态拼接。3.4 用一个小实验验证效果伪实现与观察指标在完整搬上大模型之前我建议先用一个可控的小实验来验证 Attention Residuals 是否值得推广。实验设置建议采用一个 100M 到 300M 参数的 decoder-only 模型在 1B 左右的纯文本数据上做预训练对比对照组是标准注意力。判断收益不能只看 loss必须同时记录注意力 entropy、梯度范数的分布、以及长距离依赖的任务指标比如一个简单的多文档问答或关键词定位任务。伪实现可以直接写成下面这样流程足够清晰class AttentionWithResidual(nn.Module): def __init__(self, d_model, num_heads, alpha): super().__init__() self.d_model d_model self.num_heads num_heads self.alpha alpha def forward(self, q, k, v, maskNone): B, T, C q.shape # 标准注意力分数 scores (q k.transpose(-2, -1)) / (C ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) prob torch.softmax(scores, dim-1) # 注意力残差:自身路径凸组合 identity torch.eye(T, deviceq.device).unsqueeze(0) prob (1 - self.alpha) * prob self.alpha * identity out prob v return out观察到理想情况下加了残差的模型在相同步数下 loss 下降更平滑注意力熵不会出现明显塌陷。如果这两点都成立再往大模型和长上下文场景推广不迟。4. 常见问题与排查技巧实录4.1 训练初期 loss 不降甚至微升加残差结构的模型如果初始化阶段 α 太大等于强行让模型把大量注意力放在自身位置上在数据拟合上反而是一种约束loss 下降会比标准注意力慢半拍。遇到这种情况别急着删掉残差先检查 α 的量级。经验上超过 0.2 的初始值对大多数文本任务都偏激进如果用的是 logits 残差形态还要确认偏置的缩放有没有超出合理范围通常从 0.1 或 1.0 量级起步即可。4.2 注意力权重变成“近似单位阵”有些实验会走向另一个极端把所有注意力都集中到自己身上。这说明 α 没有被正确约束或者残差项参与了反向传播而梯度主导了注意力参数更新导致网络学会了“走捷径”。排查方法是把注意力权重的对角线均值打印出来如果超过 0.5基本可以判定残差过度主导。解决方法是调小 α、加上限制残差贡献上限的正则项或者把 α 从可学习改成固定的退火策略。4.3 序列长度变化导致行为漂移Attention Residuals 中的单位阵对序列长度很敏感短序列下单位阵占比相对较高自连接的影响被放大长序列下单位阵相对稀疏影响被稀释。同一个模型被用在长短差异极大的场景时比如摘要生成和单 token 分类表现会不稳定。这是残差结构天然的尺度敏感性。实操中最好让 α 关于序列长度做归一化或者直接放弃固定单位阵改用与长度无关的游走先验。4.4 显存和推理延迟的意外上涨如果残差加在权重层面且实现方式是完整构造单位矩阵长序列下显存里会多出一份 T×T 的矩阵这是不必要的浪费。单位矩阵虽然在逻辑上是稀疏的但像上面伪代码那样的构造方法在 PyTorch 里是老老实实开了一块稠密内存。优化方法是完全绕开单位阵的构造利用注意力分数矩阵的对角元素直接调整或者使用稀疏矩阵乘法。放到推理环境里还要小心避免在多轮解码中重复构造残差矩阵提前缓存好一份即可。5. 这篇论文值得关注的下游影响与延伸思考5.1 对长上下文的直接价值长上下文场景是注意力坍缩的重灾区我现在测试一个模型在 32K 上下文里的表现时最先看的就是注意力分布有没有集中在局部或个别 token 上。Attention Residuals 如果稳定工作长上下文任务里“越远越失效”的经典瓶颈就能得到一定缓解。尤其是那些需要把信息跨多个不连续段落关联的任务一个保底的自身路径会让模型更容易维持对自身的状态一致性。当然最后还是要看论文在长序列基准上的实际收益但方向上是值得押注的。5.2 对 MoE 与多模态架构的启发残差思想一旦在注意力内部站住脚很容易外溢到其他架构。比如 MoE 模型里每个专家头的路由分布本身也依赖 top-k 选择也存在坍缩问题——某个专家永远被选中其他专家常年空转。如果把“残差保底”的思路用于路由决策给每个专家一个基础选择概率或许能缓解这种失衡。多模态对齐里文本 token 和图像 token 的注意力很容易被某一种模态主导给另一模态加一个残差路径也可能让模型更稳定地融合双通道信息。5.3 和第 2 篇的衔接还可以往哪走这篇解读的标题明确是第 1 篇说明后面大概率还有系列内容。如果之后要拆解第 2 篇我比较期待几个方向一是对多层残差链路的详细分析二是针对具体任务的一手评测数据三是和现有长期记忆结构的组合方式。从 Attention Residuals 出发完全可以接着写“如何把它和检索增强结合”“如何扩展到线性注意力”这些话题。这种工作真正有价值的点不在于提出一个完全陌生的机制而在于把一个经典思想放到了一个让人意想不到的位置上。注意力从来都不缺少复杂度缺的是那种能防退化、能稳梯度、能让大家改完还想继续改下去的简洁结构。后续工作无论朝哪个方向延伸都会有不少人跟进复现因为它足够简单又刚好戳中了实际训练里大家都遇到过、却一直没有系统解决的痛点。
注意力机制的演进已经过了堆砌结构的阶段真正有价值的改动都藏在“最朴素的直觉”里。Kimi 团队这篇关于 Attention Residuals 的论文标题看起来平淡但把“残差”和“注意力”这两个关键词放在一起懂行的人马上会精神起来。它想解决的是注意力机制里一个很少有人正面处理的问题当注意力分布逐渐走向极端模型的信息提取能力其实是被自己削弱掉的。这篇解读我会按“问题动机 → 设计空间 → 工程落地 → 踩坑排查”的顺序拆开讲适合正在做注意力变体实验、或者想在自己模型上引入稳定收益结构的研究人员和工程师参考。1. 注意力机制为什么需要“残差”先从“坍缩”说起1.1 标准注意力不是你想的那样稳定Transformer 里的自注意力数学上就三行Q 和 K 做点积除以根号 dsoftmax 成权重再加权求和 V。这个机制在绝大多数场景下表现很好但它有一个被默认接受、其实很危险的特性softmax 会把输入向量的差异放大成概率差异一旦某些 token 的注意力分数比其他 token 高出一点经过指数运算后权重会瞬间向这些 token 集中。我见过太多人调模型时碰到的诡异现象深层网络训练到一半注意力图变成了几条又细又亮的线几乎所有 token 只看跟自己最近的两三个位置或者死盯着某个语义中心词不放。这不是偶然而是注意力坍缩的典型症状。序列越长越明显因为点积分数的方差会随着参与计算的元素增多而不稳定softmax 的输出越来越接近 one-hot 分布。这种坍缩带来的后果很直接模型名义上有全局建模能力实际只会用一小截上下文长文本里的“大海捞针”任务表现一言难尽逻辑推理时也容易漏掉远处的约束条件。过去很多解法都在打补丁比如温度缩放、flash attention 里的在线归一化优化、或者给注意力分数加上相对位置偏置如 ALiBi。这些手段缓解了分布极端化但都没有把“注意力本身退化”当成一个独立的建模问题来对待。1.2 残差思想从 ResNet 到注意力模块的迁移逻辑深度学习里遇到“退化”问题最先想到的应该是残差结构。ResNet 时代的核心洞见是与其让网络硬学一个复杂的映射 H(x)不如让它学 F(x) H(x) − x再在输出端加回 x。这样一来网络不需要从头记住恒等映射只要在已有表示上做增量调制就够了。这个思路让深层网络的梯度、表达能力、训练稳定性同时受益。注意力机制其实也有对应的“退化”问题。标准注意力的输出是对 V 的加权求和如果注意力权重坍缩成 one-hot那输出就直接退化成某个位置本身的向量multi-head 的多样性也跟着没了。换句话说注意力模块的“恒等映射”应该是保留每个 token 自身的原始信息就像 ResNet 里输出端必须有一条 x 的旁路。很遗憾标准注意力里恰恰没有这条旁路每个 token 所能获取的信息完全由归一化过的权重说了算。所以 Attention Residuals 这个概念本质上是把 ResNet 的问题意识搬进了注意力计算内部在注意力得分、注意力权重、或者注意力输出这三个关键位置上显式地保留一条来自“自身”的恒等路径让注意力不会因为归一化而弄丢原始的 token 身份信息。这个迁移逻辑听起来简单但选择在哪一个位置插入残差、用什么方式保底、如何和 mask 与归一化兼容就是论文真正在做文章的地方。2. Attention Residuals 的设计空间三种可能的落地形态2.1 形态一对注意力分数logits施加结构先验第一种实现方式是在 softmax 之前的 logits 上直接加一个残差矩阵 R让注意力分数变成 QKᵀ/√d R。这个 R 可以是固定的也可以是可学习的最常见的选择是把单位阵 I 的缩放版本作为先验。这个设计想表达的逻辑是什么本质上它是在给“自己 attending 到自己”这件事一个固定的偏置加分。我把它理解成给注意力机制安装了一个保安不管你最终对全局做了多少加权组合至少你要保证不会彻底忽略掉当前位置的信息。这对深层网络特别有意义因为多层叠加之后token 的原始语义经过反复变换已经非常抽象如果没有自连接保底浅层学到的部分特征可能在深层被注意力机制直接清洗掉。这类残差与已有工作的区别值得注意。位置编码偏置比如 ALiBi是在 logits 上根据两个 token 的相对距离加一个负向偏置它的目的是约束“远距离关系”而 Attention Residuals 里加的 I 矩阵修正的是“同一位置的恒等路径”两者不冲突甚至可以叠加使用。实现层面也干净R 通常是稀疏的几乎不增加显存访问。2.2 形态二对注意力权重做凸组合给行归一化留一条恒等路第二种形态更直接在 softmax 之后把注意力权重矩阵 A 和某个残差矩阵 R 做凸组合得到新的权重 A。如果 R 取单位阵 I就是给每个位置固定分配一部分注意力给自己用纯数学语言写就是A (1 − α) · softmax(QKᵀ / √d) α · I这个公式的好处是直观α 相当于一个“自我注意力比例”控制模型保留多少基础的自身路径。你觉得 α 是 0.05 还是 0.2取决于任务里自身信息的重要程度。如果任务本身几乎完全依赖长程检索α 太大反而会拖后腿。这里要坚持一个原则加完残差之后注意力权重必须依然满足行归一化。凸组合天然满足但如果你是直接在 softmax 的结果上加上一个未缩放的单位阵那这一行的概率总和会变成 1 α后面接的因果采样和 KV Cache 处理都会出问题。常见做法是利用凸组合保证归一化依旧成立代码里一行就能实现却常常被人忽略。2.3 形态三在注意力输出上做门控残差第三种形态把残差的位置挪到了输出端即注意力机制的输出不是一个纯 V 的加权和而是原始输入 x 和注意力输出的门控混合。这一步和 Transformer 原本的 Add Norm 有相似之处但角度不同原来的 Add 是在子层外做残差连接而这里是在注意力计算内部就把原始输入作为一份成分混进去。这种设计的核心动机是增强模块的“保底能力”。深层 Transformer 里靠 LayerNorm 和残差连接外加一个注意力子层理论上信息是可以跨层流动的但实际上子层内部多次线性变换后梯度已经变得非常碎片化。如果能在子层内部再加一道旁路相当于给梯度增加了一条更短的传播通道。2.4 多头分工与层间残差链再往上走一个层次Attention Residuals 还可以从“头”和“层”两个维度展开。多头注意力原本的设计意图是让不同的头捕获不同的关系模式但坍缩发生时多个头会趋同于同一个极端分布分工就形同虚设。一种直接的扩展是给每个头设置独立的 α让一部分头偏向自我保底另一部分头专注全局检索。这种“分工式残差”引导不同头走向不同的功能角色而不是放任它们自由趋同。在层间维度早期 Transformer 层输出的注意力信息经过后续层后容易被更高阶的特征覆盖。论文标题里的“Residuals”如果放在层间视角下还能理解成跨层的注意力跳跃连接——类似把浅层的注意力分布直接送到深层去参与计算。这类做法的风险在于实现复杂度上升对推理框架的改造比较大所以通常是三种基础形态外的进阶方向。3. 从论文构思到工程实现关键细节与实操参数3.1 先线性插值还是先做归一化两种公式的等价性与差异这里有个非常容易踩坑的细节把残差加在 softmax 之前和加在 softmax 之后数学上并不总是等价。加在 softmax 之前的那条路残差矩阵会和 QKᵀ 一起参加指数运算于是残差项的影响其实是整个分数分布的函数带着复杂的非线性耦合加在 softmax 之后做凸组合则是在已经形成的概率分布上做线性混合行为直白可预期。我在实践中更推荐先想清楚“你想让残差影响模型的方式是哪种”。如果你希望残差能动态对抗 logits 方差过大导致的坍缩那放在 softmax 之前更直接如果你只想要一个稳定、易解释、可退化为标准注意力的结构那放在 softmax 之后做凸组合更省心。论文如果只讲一种实现你理解之后一定要自己把这个对照组跑一遍因为两种方式在深层网络里的表现差异往往比想象中更大。推荐把两种公式明确写成# 形态一:logits 残差 logits (q k.T) / math.sqrt(d) residual_bias prob torch.softmax(logits, dim-1) # 形态二:权重凸组合 prob (1 - alpha) * torch.softmax((q k.T) / math.sqrt(d), dim-1) alpha * identity_matrix3.2 超参数怎么设α 的初始化、退火与自适应对凸组合形态来说α 是最关键的超参数。我见到不少复现党上来就拍脑袋设 α 0.1训完说提升太小这往往是因为 α 设置不合适。按我的经验α 的初始值要跟训练阶段绑在一起。模型刚起步时注意力分布没有成型α 可以稍大一些相当于给模型一个比较强的“自我保底”先验帮助它稳定训练训练中后期各类模式逐渐成熟如果 α 还维持很大模型会变得“短视”不善于检索远程信息。一个常见做法是让 α 从 0.1 或 0.15 开始线性退火到 0.02 左右或者干脆通过一个可学习的标量来自适应。可学习的好处是省去调参但要注意对可学习标量做正则否则它容易滑到某个极端值而不是停留在最优区间。我通常还会配合观察每一层的 attention entropy。如果某几层的 entropy 明显比其他层低说明这些层已经进入坍缩状态应该给它们更大的 α这就是“按层自适应”的朴素实现。3.3 兼容性问题FlashAttention、因果掩码与 KV Cache凡是改动注意力计算的方案我都建议先考虑它跟 FlashAttention 和因果掩码兼容吗Attention Residuals 在 logits 上做文章的话天然兼容在线 softmax因为只是在分数矩阵上加一个偏置但如果在权重矩阵层面做凸组合FlashAttention 的分块计算会遇到麻烦因为我们要把分块 softmax 的结果直接和单位阵混合这破坏了它标准的 online softmax 流程需要额外遍历一次或单独处理对角块。因果掩码的影响也不容忽视。做语言模型时必须保留 causal mask那么残差单位阵 I 在对角线以下的部分其实是无效的因为那些位置被 mask 掉了。正确的姿势是让残差矩阵只在对角线及其合法范围内存在否则等于从句把未来位置的数值泄漏进来了。这个 bug 很隐蔽报错不会出现但训练效果会莫名变差。KV Cache 层面小 α 下的残差几乎不引入额外的内存开销但如果做多头独立 α缓存里就要多存一份每头的 α 向量。虽然量很小但在推理框架里这个额外的元数据会带来一点调度开销尤其对那种一次推理成千上万个请求的线上服务。工程上最好把它压进浮点常量或者量化存储别在每步解码时都动态拼接。3.4 用一个小实验验证效果伪实现与观察指标在完整搬上大模型之前我建议先用一个可控的小实验来验证 Attention Residuals 是否值得推广。实验设置建议采用一个 100M 到 300M 参数的 decoder-only 模型在 1B 左右的纯文本数据上做预训练对比对照组是标准注意力。判断收益不能只看 loss必须同时记录注意力 entropy、梯度范数的分布、以及长距离依赖的任务指标比如一个简单的多文档问答或关键词定位任务。伪实现可以直接写成下面这样流程足够清晰class AttentionWithResidual(nn.Module): def __init__(self, d_model, num_heads, alpha): super().__init__() self.d_model d_model self.num_heads num_heads self.alpha alpha def forward(self, q, k, v, maskNone): B, T, C q.shape # 标准注意力分数 scores (q k.transpose(-2, -1)) / (C ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) prob torch.softmax(scores, dim-1) # 注意力残差:自身路径凸组合 identity torch.eye(T, deviceq.device).unsqueeze(0) prob (1 - self.alpha) * prob self.alpha * identity out prob v return out观察到理想情况下加了残差的模型在相同步数下 loss 下降更平滑注意力熵不会出现明显塌陷。如果这两点都成立再往大模型和长上下文场景推广不迟。4. 常见问题与排查技巧实录4.1 训练初期 loss 不降甚至微升加残差结构的模型如果初始化阶段 α 太大等于强行让模型把大量注意力放在自身位置上在数据拟合上反而是一种约束loss 下降会比标准注意力慢半拍。遇到这种情况别急着删掉残差先检查 α 的量级。经验上超过 0.2 的初始值对大多数文本任务都偏激进如果用的是 logits 残差形态还要确认偏置的缩放有没有超出合理范围通常从 0.1 或 1.0 量级起步即可。4.2 注意力权重变成“近似单位阵”有些实验会走向另一个极端把所有注意力都集中到自己身上。这说明 α 没有被正确约束或者残差项参与了反向传播而梯度主导了注意力参数更新导致网络学会了“走捷径”。排查方法是把注意力权重的对角线均值打印出来如果超过 0.5基本可以判定残差过度主导。解决方法是调小 α、加上限制残差贡献上限的正则项或者把 α 从可学习改成固定的退火策略。4.3 序列长度变化导致行为漂移Attention Residuals 中的单位阵对序列长度很敏感短序列下单位阵占比相对较高自连接的影响被放大长序列下单位阵相对稀疏影响被稀释。同一个模型被用在长短差异极大的场景时比如摘要生成和单 token 分类表现会不稳定。这是残差结构天然的尺度敏感性。实操中最好让 α 关于序列长度做归一化或者直接放弃固定单位阵改用与长度无关的游走先验。4.4 显存和推理延迟的意外上涨如果残差加在权重层面且实现方式是完整构造单位矩阵长序列下显存里会多出一份 T×T 的矩阵这是不必要的浪费。单位矩阵虽然在逻辑上是稀疏的但像上面伪代码那样的构造方法在 PyTorch 里是老老实实开了一块稠密内存。优化方法是完全绕开单位阵的构造利用注意力分数矩阵的对角元素直接调整或者使用稀疏矩阵乘法。放到推理环境里还要小心避免在多轮解码中重复构造残差矩阵提前缓存好一份即可。5. 这篇论文值得关注的下游影响与延伸思考5.1 对长上下文的直接价值长上下文场景是注意力坍缩的重灾区我现在测试一个模型在 32K 上下文里的表现时最先看的就是注意力分布有没有集中在局部或个别 token 上。Attention Residuals 如果稳定工作长上下文任务里“越远越失效”的经典瓶颈就能得到一定缓解。尤其是那些需要把信息跨多个不连续段落关联的任务一个保底的自身路径会让模型更容易维持对自身的状态一致性。当然最后还是要看论文在长序列基准上的实际收益但方向上是值得押注的。5.2 对 MoE 与多模态架构的启发残差思想一旦在注意力内部站住脚很容易外溢到其他架构。比如 MoE 模型里每个专家头的路由分布本身也依赖 top-k 选择也存在坍缩问题——某个专家永远被选中其他专家常年空转。如果把“残差保底”的思路用于路由决策给每个专家一个基础选择概率或许能缓解这种失衡。多模态对齐里文本 token 和图像 token 的注意力很容易被某一种模态主导给另一模态加一个残差路径也可能让模型更稳定地融合双通道信息。5.3 和第 2 篇的衔接还可以往哪走这篇解读的标题明确是第 1 篇说明后面大概率还有系列内容。如果之后要拆解第 2 篇我比较期待几个方向一是对多层残差链路的详细分析二是针对具体任务的一手评测数据三是和现有长期记忆结构的组合方式。从 Attention Residuals 出发完全可以接着写“如何把它和检索增强结合”“如何扩展到线性注意力”这些话题。这种工作真正有价值的点不在于提出一个完全陌生的机制而在于把一个经典思想放到了一个让人意想不到的位置上。注意力从来都不缺少复杂度缺的是那种能防退化、能稳梯度、能让大家改完还想继续改下去的简洁结构。后续工作无论朝哪个方向延伸都会有不少人跟进复现因为它足够简单又刚好戳中了实际训练里大家都遇到过、却一直没有系统解决的痛点。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
