Attention Residuals

问题动机

现代LLM普遍采用PreNorm（Pre-LayerNorm）结构的残差连接，但存在一个根本性缺陷：所有层的输出以固定单位权重（权重=1）累加到残差路径。这种均匀聚合导致两个问题：

1. 隐藏状态无控增长：深层网络的层输出幅度逐层累积，导致数值不稳定
2. 层贡献稀释：每层的独立贡献被后续层的累加所稀释，信息传递效率降低

标准公式为： $\mathbf{y} = \mathbf{x} + f_L(\mathbf{x})$ $\mathbf{x}_{l+1} = \text{PreNorm}(\mathbf{y})$

方法核心：Attention Residuals

AttnRes的核心思想是用**softmax attention**替代固定累加，让每层自适应地选择性地聚合之前的层输出：

$\mathbf{y}_l = \alpha_l \cdot \mathbf{x}_l + (1-\alpha_l) \cdot f_l(\mathbf{x}_l)$ 其中权重 $\alpha_l = \text{softmax}(\mathbf{q}_l \mathbf{K}_{<l})$ 由当前层查询与之前层键的交互决定。

更精确的AttnRes形式： $\mathbf{z}_l = \text{Attention}(\mathbf{Q}=\mathbf{x}_l, \mathbf{K}=\mathbf{h}_{<l}, \mathbf{V}=\mathbf{h}_{<l})$ $\mathbf{h}_l = \lambda_l \cdot \mathbf{x}_l + (1-\lambda_l) \cdot \mathbf{z}_l$

其中 $\lambda_l$ 是可学习的标量，控制残差与注意力的平衡。

关键特性：

• 输入依赖：权重 $\alpha_l$ 随输入内容变化
• 深度选择：每层可独立决定保留多少原始信号、聚合多少深层信息
• 可学习门控：$\lambda_l$ 提供额外的层级别控制

Block AttnRes工程优化

直接对所有前面层做attention会导致 $O(L^2)$ 的内存和通信开销。Block AttnRes通过层级分块解决：

1. 将 $L$ 层划分为 $G$ 个块，每块 $B = L/G$ 层
2. 每块维护一个块级表示 $\mathbf{b}_g$，该块最后一层的输出
3. Attention在块级表示上进行：$\mathbf{z}_l = \text{Attention}(\mathbf{x}_l, \{\mathbf{b}_0, ..., \mathbf{b}_{g-1}\})$
4. 块内仍使用标准残差连接

复杂度从 $O(L^2)$ 降至 $O(B \cdot G + B^2)$，实际应用中 $B \ll L$ 时显著降低。

实现优化：

• Cache-based pipeline通信：跨节点通信块级表示而非所有层
• 两阶段计算：先计算块内，再聚合块间

验证结果

• Scaling law：跨0.1B到几十B参数规模一致有效
• Ablation：验证输入依赖的深度选择确实带来收益
• Kimi Linear (48B/3B)：1.4T tokens预训练，输出幅度和梯度分布更均匀，下游任务全面提升