[Kimi 团队] Attention Residuals：打破固定残差枷锁，让大模型在“深度空间”学会注意力

TL;DR

在现代 LLM 中，残差连接（Residual Connections）通常被视为“理所当然”的单位权重累加。然而，Kimi 团队在最新技术报告中指出：这种缺乏选择性的累加是深层模型“稀释效应”的元凶。他们提出的 Attention Residuals (AttnRes) 将残差连接从“无脑加法”升级为“内容相关的注意力检索”，不仅在 Scaling Law 上表现出 1.25 倍的计算效率优势，更在复杂推理任务中显著超越了传统的 PreNorm 架构。

背景：被忽视的“残差稀释”痛点

在 Transformer 的标准公式 $h_l=h_{l-1}+f_{l-1}(h_{l-1})$ 中，每一层对最终状态的贡献权重都是固定的“1”。

1. 层稀释（Layer Dilution）：随着模型变深，隐状态幅值线性增长。早期层的信息被深层输出不断掩埋，模型难以“回头看”。
2. 缺乏选择性：无论当前是 Attention 还是 MLP 层，它们接收到的都是完全相同的聚合状态，无法针对性地从前序层中提取特定信息。

这种现象非常像 Transformer 出现之前的序列建模——那时的 RNN 依靠这种递推公式通过时间传递信息，而现在，AttnRes 决定对“深度”这一维度也应用注意力机制。

核心方法：深度维度的 Softmax 注意力

AttnRes 的直觉非常优雅：既然每一层都在产生新的 Representation，为什么不让后续层通过一个学到的 Pseudo-query ($w_l$) 来对前面所有层的输出做一次 Softmax 检索呢？

$h_l={\sum }_{i=0}^{l-1}{\alpha }_{iol}\cdot v_i$

其中 $\alpha$ 是基于内容的相关性权重。

解决落地难题：Block AttnRes

虽然全量注意力（Full AttnRes）在小规模测试中有效，但在大规模训练（Pipeline Parallelism）下，每一层都要把所有前序状态传给下一个 GPU 节点，通信开销会爆炸。

• 分块聚合：将 $L$ 层划分为 $N$ 个块。
• 块间检索：内层仍然使用常规残差，但跨块时使用注意力检索。
• 极简开销：仅仅增加了一个学到的 $d$ 维向量作为 Query 和一个 RMSNorm，参数量增加几乎可以忽略不计。

上图展示了学习到的深度注意力权重分布。我们可以清晰地看到明显的“对角线主导（局部性）”以及对“第 0 层（Embedding）”的持续关注，这说明模型学会了动态平衡基础特征与深层抽象。

效率优化：两阶段推理策略

为了不牺牲推理速度，作者设计了 Two-phase Computation：

1. Phase 1（并行）：一次性读入所有块的 KV，完成块间注意力的计算。
2. Phase 2（顺序）：在块内通过 Online Softmax 逐层更新部分和，并与 Phase 1 的结果融合。 这种设计将内存 I/O 从 $O(L)$ 降到了 $O(S+N)$，实测推理延迟增加小于 2%。

实验战绩：推理能力的飞跃

在 Kimi Linear 48B（3B 激活参数）模型上的 1.4T token 预训练显示，AttnRes 的提升主要集中在硬核推理领域：

• GPQA-Diamond（研究生水平难题）：从 36.9 暴涨至 44.4 (+7.5)。
• MATH：从 53.5 提升至 57.1 (+3.6)。
• HumanEval（代码生成）：提升 +3.1。

Scaling Law 曲线清楚地显示：在相同的损失函数目标下，AttnRes 持续领先基线，相当于获得了额外的算力加成。

深度洞察：为什么有效？

传统的 PreNorm 在训练后期，梯度往往会集中在最前面的几层。AttnRes 通过学到的梯度路由：

1. 梯度分布更均匀：Softmax 的竞争机制让每一层都能获得合理的梯度信号。
2. 幅值更稳定：观察数值动态发现，AttnRes 能有效抑制隐状态的爆炸性增长，呈现出健康的周期性波动。

总结与未来启示

Attention Residuals 的成功标志着“统一建模时间与深度”进入了实用阶段。它告诉我们：

• 深度不仅是叠加：深度也是一种可以被检索的存储空间。
• 架构可塑性：通过 Block 机制，我们可以在计算成本和模型表达能力之间找到新的 Pareto 最优解。

虽然目前 $O(Ld)$ 的通信仍是 Full AttnRes 的瓶颈，但随着未来高速互联硬件的发展，或许未来的大模型每一层都能动态“瞬移”并提取模型任何位置的特征。