用注意力机制替代 Transformer 的标准残差连接——让每一层选择性地聚合所有前层表示,而非简单均匀叠加。一个即插即用的改动,带来跨越算力规模的持续收益。
PreNorm Transformer 每一层以固定权重 1 叠加前层输出。随着网络加深,这种均匀累积导致层贡献被稀释,且隐藏状态的幅值随深度线性增长——无界。
每层输出不断叠加,隐藏状态的模长随深度线性增长,导致训练不稳定——这是 PreNorm 架构的已知缺陷。
网络越深,单层输出对最终表示的相对影响越小,导致浅层梯度信号衰弱,深层模型收益递减。
聚合权重与输入内容完全无关,无论当前 token 是代码、数学还是文本,各层贡献始终等权。
用 softmax 注意力替代固定叠加。每层通过一个学习到的伪查询向量,对所有前层表示进行内容感知的加权聚合,权重随输入动态变化。
每层可直接访问任意更早层的表示,无需信息逐层传递,消除深层网络中的信息瓶颈与梯度消失问题。
注意力权重 α 由输入决定。处理代码时可更多关注语法层,处理数学时可更多关注计算层,真正做到选择性聚合。
Full AttnRes 需要 O(L·d) 的内存缓存所有层输出。Block AttnRes 将层分组为 N 个块,块内用标准残差,仅在块级别施加跨块注意力,将内存降至 O(N·d)。约 8 个块即可恢复绝大部分收益。
N 为块数,远小于层数 L。实验中 N≈8 已足够,相比 Full AttnRes 节省 6× 以上的内存开销。
仅需替换残差计算逻辑,无需改动整体架构,边际计算开销极低,易于集成到现有训练框架。
在 Kimi Linear 48B MoE 模型(3B 激活参数)上验证:约 8 个块是性能与效率的最佳平衡点。
在 Kimi Linear 48B / 1.4T tokens 上验证,AttnRes 在全部 9 项基准中均优于基线,多步推理(+7.5 GPQA-Diamond)和代码生成(+3.1 HumanEval)增益尤为显著。
| Benchmark | Baseline | AttnRes | Δ |
|---|---|---|---|
| 通用能力 | |||
| MMLU | 73.5 | 74.6 | |
| GPQA-Diamond | 36.9 | 44.4 | |
| BBH | 76.3 | 78.0 | |
| TriviaQA | 69.9 | 71.8 | |
| 数学与代码 | |||
| MATH | 53.5 | 57.1 | |
| HumanEval | 59.1 | 62.2 | |
| MBPP | 72.0 | 73.9 | |
| 中文能力 | |||
| CMMLU | 82.0 | 82.9 | |
| C-Eval | 79.6 | 82.5 | |
AttnRes 有效缓解了 PreNorm 的幅值爆炸:隐藏状态模长在深度方向保持有界,各层梯度范数分布更均匀,训练过程更加稳定可控。