Attention Residuals

论文笔记：Attention Residuals

元信息

一句话总结

用 Softmax attention 替代固定权重的 残差连接，让每层可选择性聚合之前所有层的输出，缓解 PreNorm dilution 问题

核心贡献

Attention Residuals (AttnRes): 将残差连接的固定累加替换为深度维度上的 learned Softmax attention，每层通过一个可学习 pseudo-query 向量选择性聚合所有前序层输出

Block AttnRes: 将层分组为 $N$ 个 block，跨 block 用 attention、block 内用标准求和，将内存和通信开销从 $O(Ld)$ 降到 $O(Nd)$

系统工程优化: cross-stage caching 消除 pipeline parallelism 下的冗余传输，two-phase 推理策略使延迟开销 <2%

问题背景

要解决的问题

标准 残差连接 以固定权重 1 累加所有前序层输出，导致三个问题：

1. 无选择性访问：不同类型的层（attention vs MLP）接收相同的聚合状态
2. 不可逆信息丢失：聚合后的信息无法选择性恢复
3. 输出幅度增长：深层必须学习越来越大的输出来在累积残差中保持影响力

现有方法的局限

PreNorm 导致隐状态幅度随深度以 $O(L)$ 增长，逐渐稀释每层的相对贡献（PreNorm dilution）

Highway Network 引入 element-wise 门控，但仍然只能访问前一层的压缩状态 $\boldsymbol{h}_{l-1}$

DenseFormer 赋予每层对所有前序输出的访问，但使用固定的、输入无关的标量系数

mHC 通过 $m$ 条并行流引入输入依赖性，但混合矩阵复杂度高，I/O 开销大（$34d$ vs AttnRes 的 $5.5d$）

本文的动机

观察到深度维度上的残差累加与序列维度上的 RNN 递推存在对偶关系

Transformer 用 attention 替代 RNN 的线性递推解决了序列建模问题；类比地，可以用 attention 替代残差连接的固定累加来解决深度维度上的信息聚合问题

网络深度 $L<1000$ 远小于序列长度，depth-wise attention 的 $O(L^2)$ 开销完全可接受

方法详解

模型架构

AttnRes 在标准 Transformer 架构上仅修改残差连接机制：

• 架构基底: Kimi Linear / DeepSeek-V3 的 MoE Transformer（KDA + MLA + MoE FFN）
• 新增参数: 每层仅增加一个 RMSNorm 和一个 $d$ 维 pseudo-query 向量 $\boldsymbol{w}_l$，参数量可忽略不计
• 初始化: 所有 pseudo-query 初始化为零，使初始 attention 权重均匀，等价于标准残差连接
• 两种变体: Full AttnRes（attention over 所有前序层）和 Block AttnRes（分组 attention）

核心模块

模块1: Full Attention Residuals

设计动机: 让每层通过 Softmax attention 选择性聚合所有前序层的输出

具体实现:

• 每层 $l$ 有一个可学习的 pseudo-query $\boldsymbol{q}_l = \boldsymbol{w}_l \in \mathbb{R}^d$
• key 和 value 来自所有前序层输出：$\boldsymbol{k}_i = \boldsymbol{v}_i = f_i(\boldsymbol{h}_i)$（$i \geq 1$），embedding 为 $\boldsymbol{v}_0 = \boldsymbol{h}_1$
• 使用 $\phi(\boldsymbol{q}, \boldsymbol{k}) = \exp(\boldsymbol{q}^\top \text{RMSNorm}(\boldsymbol{k}))$ 作为 kernel function
• 输入到第 $l$ 层为 $\boldsymbol{h}_l = \sum_{i=0}^{l-1} \alpha_{i \to l} \cdot \boldsymbol{v}_i$
• 内存开销 $O(Ld)$，计算 $O(L^2 d)$

模块2: Block Attention Residuals

设计动机: Full AttnRes 在 pipeline parallelism 下每层输出都要跨 stage 传输，通信开销 $O(Ld)$ 不可接受

具体实现:

• 将 $L$ 层分为 $N$ 个 block，每 block $S = L/N$ 层
• block 内: 标准残差求和 $\boldsymbol{b}_n = \sum_{j \in \mathcal{B}_n} f_j(\boldsymbol{h}_j)$
• block 间: 对 $N$ 个 block 级别表示做 Softmax attention
• 当前 block 内的层还额外 attend 到不断累积的 intra-block partial sum $\boldsymbol{b}_n^i$
• 内存和通信降至 $O(Nd)$，计算 $O(N^2)$
• 经验发现 $N \approx 8$ 即可恢复 Full AttnRes 的大部分收益

模块3: Two-Phase Computation（推理优化）

Phase 1: 批量计算 block 内所有 $S$ 层的 inter-block attention（共享 KV cache），将内存访问从 $S$ 次降到 1 次

Phase 2: 顺序计算 intra-block attention，用 online softmax 与 Phase 1 结果合并

总 per-layer I/O 仅 $(\frac{N}{2}+5)d$（典型值 $5.5d$），vs 标准残差的 $3d$

关键公式

公式1: 标准残差更新

含义: 标准残差连接以固定权重 1 累加，展开后 $\boldsymbol{h}_l = \boldsymbol{h}_1 + \sum_{i=1}^{l-1} f_i(\boldsymbol{h}_i)$

符号说明:

• $\boldsymbol{h}_l \in \mathbb{R}^d$: 第 $l$ 层的隐状态
• $f_l$: 第 $l$ 层的变换函数（self-attention 或 MLP）

公式2: Attention Residuals 核心公式

含义: 用 learned attention 权重 $\alpha_{i \to l}$ 替代固定权重 1，实现选择性深度聚合

符号说明:

• $\alpha_{i \to l}$: 层 $l$ 对层 $i$ 输出的 attention 权重，$\sum_{i=0}^{l-1} \alpha_{i \to l} = 1$
• $\boldsymbol{h}_1$: token embedding

公式3: Attention 权重计算

含义: 通过 kernel function $\phi$ 计算归一化的 attention 权重

符号说明:

• $\phi(\boldsymbol{q}, \boldsymbol{k}) = \exp(\boldsymbol{q}^\top \text{RMSNorm}(\boldsymbol{k}))$: 保证非负的 kernel function
• $\boldsymbol{q}_l = \boldsymbol{w}_l$: 可学习的 pseudo-query 向量
• $\boldsymbol{k}_i = \boldsymbol{v}_i$: key/value 共享，来自前序层输出

公式4: Full AttnRes 层输入

含义: Full AttnRes 的简洁形式，每层输入是所有前序层输出的 attention-weighted sum

符号说明:

• $\boldsymbol{v}_0 = \boldsymbol{h}_1$ (embedding), $\boldsymbol{v}_i = f_i(\boldsymbol{h}_i)$ for $1 \leq i \leq l-1$

公式5: Block 内累加

含义: 每个 block 内的层输出通过标准求和压缩为单一表示

符号说明:

• $\mathcal{B}_n$: 第 $n$ 个 block 包含的层索引集合
• $\boldsymbol{b}_n^i$: 前 $i$ 层的 partial sum

公式6: Block AttnRes 值矩阵

含义: block 首层 attend 到所有已完成 block；后续层额外 attend 到当前 block 的 partial sum

公式7: 朴素通信开销

含义: 不做优化时每次 stage 转换要传输所有累积的 block representations

符号说明:

• $C = PV$: 总 chunk 数（$P$ 物理 stage $\times$ $V$ 虚拟 stage）
• $N_p$: 每个物理 stage 平均产生的 block 数

公式8: 缓存优化通信

含义: cross-stage caching 使峰值通信从 $O(C)$ 降到 $O(P)$，实现 $V\times$ 的改善

公式9: Sequence-Depth 对偶——TTT 递推

含义: Test-Time Training 将序列递推形式化为梯度下降，展示了时间-深度对偶性

符号说明:

• $\mathbf{W}_t$: 第 $t$ 步的模型状态
• $\ell$: 自监督损失

公式10: mHC 权重展开

含义: mHC 的 depth mixing 权重可以展开为 structured matrix 形式，揭示其 $m$-semiseparable 结构

符号说明:

• $\mathbf{A}_{i \to j}^\times := \prod_{k=i+1}^{j} \mathbf{A}_k$: 累积转移矩阵
• $\boldsymbol{\alpha}_l, \boldsymbol{\beta}_i \in \mathbb{R}^m$: mixing 向量

公式11: Scaling Law 拟合

Baseline: $\mathcal{L} = 1.891 \times C^{-0.057}$

Block AttnRes: $\mathcal{L} = 1.870 \times C^{-0.058}$

Full AttnRes: $\mathcal{L} = 1.865 \times C^{-0.057}$

含义: 三种方法 slope 相近，但 AttnRes 一致性地在整个 compute range 上取得更低 loss

关键图表

Figure 1: Overview / 三种残差连接对比

{:width 600}

说明: (a) Standard Residuals: 固定权重 1 的加性累加；(b) Full AttnRes: 每层通过 learned attention 选择性聚合所有前序层输出；(c) Block AttnRes: 层分组为 block，block 内求和、block 间 attention，内存从 $O(Ld)$ 降到 $O(Nd)$

Figure 2: PyTorch 伪代码 / Block AttnRes 实现

说明: block_attn_res 函数接收 block representations 列表和 intra-block partial sum，通过 torch.einsum 计算 pseudo-query 与 RMSNorm 归一化后的 keys 之间的 logits，再用 Softmax 加权聚合。forward 函数展示了完整的层内流程：先 apply AttnRes → attention → AttnRes → MLP，在 block 边界处将 partial_block 加入 blocks 列表

Figure 3: Pipeline 通信优化

说明: 4 个物理 rank、2 个虚拟 stage 的 cache-based pipeline 通信示例。每个 rank 缓存之前收到的 block representations（hatched boxes），stage 转换时仅传输增量 block（如 $+[\boldsymbol{b}_1, \boldsymbol{b}_2]$），消除了 6 次冗余传输

Figure 4: Scaling Law 曲线

{:width 600}

说明: Full 和 Block AttnRes 在所有 compute budget 下都一致超越 baseline。在 5.6 PFLOP/s-days 时，Block AttnRes 达到 1.692 vs Baseline 的 1.714，相当于 1.25x compute advantage。Block 与 Full 的差距随 scale 缩小，最大规模时仅差 0.001

Figure 5: 训练动态对比

{:width 600}

说明: (a) Validation loss: AttnRes 全程更低，decay 阶段差距拉大；(b) Output magnitude: Baseline 随深度单调增长（PreNorm dilution），Block AttnRes 呈有界周期性模式（block 边界重置累加）；(c) Gradient magnitude: Baseline 梯度在浅层异常大，AttnRes 的 Softmax 竞争机制使梯度分布更均匀

Figure 6: Block Size 消融

说明: 16层模型上 Block Size $S$ 的影响。$S=1$ 即 Full AttnRes (loss 1.737)，$S=2,4,8$ 都接近 1.746，$S=16$ (=1 block) 退化为 1.757 接近 baseline (1.766)。$N \approx 8$ 即可恢复大部分收益

Figure 7: 架构搜索热力图

说明: 固定 compute ($6.5 \times 10^{19}$ FLOPs) 和参数量 ($2.3 \times 10^8$) 下的 25 种 $(d_\text{model}/L_b, H/L_b)$ 配置。关键发现: Baseline 最优点在 $d_\text{model}/L_b \approx 60$ (1.847)，AttnRes 将最优点移至 $d_\text{model}/L_b \approx 45$ (1.802)，即 AttnRes 偏好更深更窄的网络

Figure 8: Attention 权重热力图

说明: 16-head 模型的 depth-wise attention 权重分布。三个关键观察: (1) 保持局部性: 对角线主导，每层主要 attend 前一层；(2) 层特化: embedding ($\boldsymbol{h}_1$) 在整个深度上保持显著权重，pre-MLP 输入更依赖近邻，pre-attention 输入保持更广的 receptive field；(3) Block 保持结构: Full→Block 压缩后关键模式全部保留，block 压缩起到隐式正则化作用

Figure 9: Depth Mixing Matrix 对比

说明: 四种残差变体的 depth mixing matrix $\mathbf{M}$（$L=4$）。Standard residual: 全 1 下三角；Highway Network: 1-semiseparable（累积门控积）；mHC: $m$-semiseparable（learned 转移矩阵）；Full AttnRes: dense rank-$L$（每个元素是独立的 $\phi(\boldsymbol{w}_l, \boldsymbol{k}_i)$ score）；Block AttnRes: 同 block 内共享 key/value，rank 介于 $N$ 和 $N+S$ 之间

Table 1: Memory I/O 对比

说明: Block AttnRes 的 per-layer I/O 仅 $5.5d$，远低于 mHC 的 $34d$，接近标准残差的 $3d$

Table 2: Scaling Law 配置与结果

说明: 所有 scale 下 AttnRes 一致超越 baseline，Block AttnRes 与 Full AttnRes 差距随 scale 缩小

Table 3: 48B 模型下游 Benchmark 结果

说明: AttnRes 在所有 benchmark 上匹配或超越 baseline。多步推理任务收益最显著: GPQA-Diamond +7.5, Math +3.6, HumanEval +3.1，支持”改善 depth-wise 信息流有利于组合性任务”的假说

Table 4: 消融实验（16层模型）

关键发现: (1) Softmax 优于 sigmoid（竞争性归一化更好）；(2) 单头优于多头（最优 depth mixture 在通道间高度一致）；(3) RMSNorm on keys 是关键组件，防止大幅度输出主导 attention

Table 5: 残差更新机制统一对比

说明: AttnRes 是唯一同时满足 dynamic weight + cross-layer access + softmax attention 的方法。标准残差对应 depth-wise linear attention，AttnRes 对应 depth-wise softmax attention

实验

数据集

实现细节

架构: Kimi Linear 48B total / 3B activated（27 Transformer blocks，54 层）

Block 配置: 6 layers/block → 9 blocks + embedding = 10 depth-wise sources

优化器: Muon optimizer

学习率: WSD (Warmup-Stable-Decay) schedule

Batch Size: 8M tokens (global)

Context Length: 4096 → 32K（逐步扩展）

预训练: Phase 1: 1T tokens WSD → Phase 2: ~400B tokens mid-training (Moonlight recipe)

位置编码: KDA 用 RoPE，MLA 用 NoPE，context extension 无需 YaRN

可视化结果

Attention 权重热力图揭示了三种模式：局部性（对角线）、embedding persistence（第 0 列）、learned skip connections（off-diagonal concentrations）

Block 压缩保留了 Full AttnRes 的核心信息通路，起到隐式正则化作用

批判性思考

优点

理论优雅: 将残差连接与 RNN/attention 的时间-深度对偶关系形式化，用 structured matrix 统一现有方法

工程实用: Block AttnRes 每层仅增加 1 个 $d$ 维向量和 1 个 RMSNorm，推理延迟开销 <2%

系统化消融: 每个设计选择（softmax vs sigmoid、单头 vs 多头、RMSNorm on keys）都有充分的消融支持

规模验证: 从 194M 到 48B 参数，scaling law 拟合到 1.4T token 预训练，验证充分

缓解 PreNorm dilution: 有界的输出幅度和均匀的梯度分布是实质性的训练稳定性改善

局限性

仅验证 MoE 架构: 所有实验基于 Kimi Linear (MoE+KDA+MLA)，未验证 dense Transformer (GPT/LLaMA style)

推理仍有额外开销: 虽然 <2%，但在对延迟极度敏感的场景仍不可忽视

Input-dependent query 更好但被放弃: loss 1.731 vs 1.737，因推理需序列依赖而未采用，存在改进空间

Block 边界硬性固定: block 划分是静态等分，未探索自适应分组

缺乏长上下文评估: 论文未报告 NIAH 或 LongBench 等长上下文 benchmark

潜在改进方向

Linear attention variant: 用 depth-wise linear attention (如 GLA/RetNet style) 替代 full softmax，进一步降低开销

Adaptive blocking: 根据层间相似度动态决定 block 边界

与其他 normalization 结合: 探索 PostNorm / HybridNorm 下 AttnRes 的效果

Dense model 验证: 在 LLaMA-style dense Transformer 上验证通用性

Input-dependent query 的高效实现: 设计不需要序列依赖的 input-dependent query 方案

可复现性评估

• 代码开源 (GitHub)
• 预训练模型（未提供）
• 训练细节完整（超参数、schedule、数据量都有详细说明）
• 数据集可获取（Kimi 内部数据）

关联笔记

基于

残差连接: He et al. 2015，本文的改进目标

PreNorm: 导致 dilution 问题的 normalization 范式

Highway Network: element-wise gating 的先驱，AttnRes 从 gating 推广到 cross-layer attention

对比

DenseFormer: cross-layer access 但 input-independent 标量系数

mHC: $m$ 并行流 + learned mixing matrix，I/O 开销高

ReZero: 可学习标量但仅访问前一层

LayerScale: 可学习 element-wise 系数但仅访问前一层

方法相关

Softmax: AttnRes 的核心注意力归一化

RMSNorm: 防止大幅度层主导 attention 权重的关键组件

online softmax: two-phase 推理中合并 inter/intra-block 结果

pipeline parallelism: Block AttnRes 的工程动机

MoE: 实验采用的模型架构

Kimi Linear: 48B 实验的基础架构

硬件/数据相关

neural scaling law: 验证 AttnRes 在不同 compute budget 下的一致收益

速查卡片

Attention Residuals

• 核心: 用 depth-wise softmax attention 替代固定残差累加，让每层选择性聚合前序层输出
• 方法: 每层一个 $d$ 维 pseudo-query → softmax over 前序层 → weighted sum；Block 变体分组压缩至 $O(Nd)$
• 结果: 1.25x compute advantage，48B 模型 GPQA +7.5, Math +3.6, HumanEval +3.1
• 代码: https://github.com/MoonshotAI/Attention-Residuals

笔记创建时间: 2026-03-20