Attention Residuals：用注意力重塑深度维度的信息聚合 Block Attention Residuals（块注意力残差）架构交互演示
论文：Attention Residuals · Kimi Team (Moonshot AI) · 2026-03 · arXiv:2603.15031 · GitHub ｜ 本页面为单文件交互演示；公式排版经 KaTeX CDN 加载，失败时保留原始 LaTeX 文本
速览 1 · 问题 2 · Full AttnRes 3 · Block AttnRes ★ 4 · 两阶段计算 5 · 效果 动画速度 0.5× 1× 2× 4×
0一分钟速览
核心洞察（时间–深度对偶）：残差连接沿深度压缩信息的方式，与 RNN 沿时间压缩信息的方式同构。 Transformer 曾用注意力取代序列上的递归；AttnRes 把同样的变革搬到深度维度：用 softmax 注意力取代固定的残差累加。
问题
PreNorm 深度稀释
标准残差 把所有层输出以固定权重 1 累加，隐藏态幅度随深度 O(L) 增长，早期层贡献被不断稀释，且无法按内容选择性读取。
方案
Attention Residuals
每层配备一个可学习的伪查询 wₗ（零初始化），对所有前序层输出做 softmax 注意力加权聚合：，权重随输入内容自适应。
落地
Block AttnRes
L 层分成 N 块：块内仍是普通残差（部分和），块间只对 N 个块表示做注意力。内存/通信从 O(L·d) 降到 O(N·d)，N≈8 即收回大部分收益。
1问题：标准残差连接的“深度稀释”
讲解要点：标准 PreNorm 残差提供了有效的梯度高速公路，但它沿深度采用固定、统一的聚合规则—— 每层输出都以固定权重 1加入主干。展开递推式可以看到：深层收到的，是所有前序层输出的等权求和。 随着层数增加，主干幅度线性膨胀，每一层（尤其早期层）的相对贡献被持续稀释。
▶ 播放 ⏭ 单步 ↺ 重置
**后果：**权重恒为 1，无法按内容选择；主干幅度 O(L) 膨胀，早期层贡献被持续稀释。 既有研究观察到，在损失有限的情况下可以裁剪非平凡比例的层 —— 等权聚合存在浪费（PreNorm dilution）。层可裁剪现象见论文所引研究（引言 [11]）。
词元嵌入 h₁ 注意力层输出 MLP/MoE 层输出 残差主干 h（等权累加）
2核心思想：Full Attention Residuals
讲解要点：既然问题出在“固定等权”，就把它换成可学习、随输入变化的 softmax 注意力。 每层 l 只有一个可学习向量——伪查询 wₗ（与输入计算解耦，并零初始化：训练开始时各源权重均匀，但此后注意力权重仍依赖输入，因为键来自词元的真实中间表示）。 键经 RMSNorm，避免大幅度表示仅因尺度支配权重；值仍是各层的实际表示。 下方热力图逐行演示与论文 Figure 8 报告一致的定性规律（对角占优、词元嵌入保持非平凡权重、偶尔出现长程跳连、Pre-Attn 与 Pre-MLP 存在不同专门化模式），但不是该图的原始测量数值。
**数据声明：**示意教学权重，并非论文 Figure 8 的原始测量数值。
▶ 播放 ⏭ 单步 ↺ 重置
⇆ 横向滑动查看完整图
注意力权重 α（颜色越深越大） 源 0 = 词元嵌入 源 i = 第 i 层输出 fᵢ
3Block Attention Residuals：可扩展的实用形态 ★
讲解要点：Full AttnRes 需要保存并传输全部 L 个层输出（O(L·d)），在大规模流水线并行训练下太贵。 Block AttnRes 把 L 层分成 N 块：块内用普通残差累加成“部分和” bₙⁱ，块结束时封存为块表示 bₙ； 块间每层只对 N 个块表示（+ 当前部分和）做注意力。内存与通信从 O(L·d) 降到 O(N·d)；缩放律实验中 N≈8 即可收回大部分收益。 下方以 L=8 层、N=4 块、每块 S=2 层为例逐层演示。
▶ 播放 ⏭ 单步 ↺ 重置
⇆ 横向滑动查看完整图
词元嵌入 b₀ = h₁ 注意力层 MLP/MoE 层 块内部分和 bₙⁱ（普通残差） 块表示缓存（O(N·d)） 注意力权重 α
4高效实现：两阶段计算策略（推理）
讲解要点：伪查询 wₗ 是可学习参数，与输入、前向计算完全解耦——因此不必等层逐个算完。 Phase 1 把一个块内全部 S 个伪查询打包成一次矩阵乘，对块表示缓存做一次批量注意力（块表示只读 1 遍而非 S 遍）； Phase 2 在块内顺序推进：第 1 层直接使用 Phase 1 的输出；第 i≥2 层再对“部分和”做一次小注意力，并用在线 softmax与 Phase 1 的结果精确合并。 为展示更一般的情形，本节以块内 S=4 层、缓存中已有 b₀~b₃ 共 4 个表示为例（与上节例子相互独立）。
▶ 播放 ⏭ 单步 ↺ 重置
I/O 与延迟（论文事实）：
论文典型配置 L=128、N=8、S=16：残差机制 I/O 为 (N/S + 5)d = 5.5d / token / layer；标准残差为 3d；mHC（m=4）为 34d。
典型工作负载下，端到端推理延迟开销低于 2%。
论文 §4.1 报告，流水线并行下 Block AttnRes 的端到端训练开销实测低于 4%（不开流水线时可忽略）。
注：本节所称“缓存”指 Block AttnRes 的块表示缓存（在深度注意力中充当 K、V），与 Transformer 序列注意力的 KV cache 无关。
块表示缓存（跨层共享的深度注意力 K、V，非序列 KV cache） 伪查询 wₗ（参数，提前已知） 批量矩阵乘 块内部分和 在线 softmax 合并
5效果：有界的幅度、更加均匀的梯度、全面的收益
**讲解要点：**左图——基线模型各层输出幅度随深度单调膨胀（PreNorm 稀释）， Block AttnRes 在块边界通过注意力“重置”累加，幅度呈有界的周期形态；右图——基线梯度集中在最浅层， AttnRes 的 softmax 竞争让各源争夺概率质量，梯度沿深度分布显著更加均匀。下图仅用于讲解论文 Figure 5 所述趋势，按 27 个 Transformer-block 位置展示。 最终 Kimi Linear 模型为 48B 总参数 / 3B 激活参数，54 个子层（每块 6 个子层，共 9 块，另含词元嵌入），训练使用 1.4T tokens。 论文 Figure 5(a) 还报告了 Block AttnRes 的验证损失全程低于基线，且在衰减阶段差距拉大（训练动态对比基于 1T tokens）。
**数据声明：**依据论文 Figure 5 所述趋势绘制的定性示意图，并非论文原始数值序列。
▶ 播放 ⏭ 单步 ↺ 重置
⇆ 横向滑动查看完整图
1.25×
计算优势：相同验证损失下，基线所需训练计算量约为 Block AttnRes 的 1.25×（scaling law）
< 2%
推理延迟开销（典型负载）
N ≈ 8
块数即收回 Full AttnRes 大部分收益
+7.5
GPQA-Diamond（48B/3B MoE, 1.4T tokens）；HumanEval +3.1，MMLU +1.1，MMLU-Pro 持平，15 个报告基准全部持平或提升
注：以上数字为论文报告的结果，并非由上方示意曲线计算得出。
基线（标准 PreNorm 残差） Block AttnRes
依据论文 Attention Residuals（Kimi Team, Moonshot AI, arXiv:2603.15031）整理制作，动画中的注意力权重为教学演示用的示意数据。
参考：arXiv 论文 · 官方仓库