SendSoon
企业事务性邮件 & AI Agent
AI 大模型精选文章

Attention Residuals

#Attention Residuals#AttnRes#Transformer#Residual Connections#PreNorm

arXiv。 arXiv:2603.15031(Submitted 16 March 2026,v1)

摘要

Residual connections with PreNorm are standard in modern LLMs, yet they accumulate all layer outputs with fixed unit weights. This uniform aggregation causes uncontrolled hidden-state growth with depth, progressively diluting each layer's contribution. We propose Attention Residuals (AttnRes), which replaces this fixed accumulation with softmax attention over preceding layer outputs, allowing each layer to selectively aggregate earlier representations with learned, input-dependent weights. To address the memory and communication overhead of attending over all preceding layer outputs for large-scale model training, we introduce Block AttnRes, which partitions layers into blocks and attends over block-level representations, reducing the memory footprint while preserving most of the gains of full AttnRes. Combined with cache-based pipeline communication and a two-phase computation strategy, Block AttnRes becomes a practical drop-in replacement for standard residual connections with minimal overhead. Scaling law experiments confirm that the improvement is consistent across model sizes, and ablations validate the benefit of content-dependent depth-wise selection. We further integrate AttnRes into the Kimi Linear architecture (48B total / 3B activated parameters) and pre-train on 1.4T tokens, where AttnRes mitigates PreNorm dilution, yielding more uniform output magnitudes and gradient distribution across depth, and improves downstream performance across all evaluated tasks.

摘要将困难定位于 PreNorm 残差设计:各层输出以单位权重累加,隐藏状态幅值随深度增长,新一层贡献被稀释。AttnRes 用对先前层输出的 Softmax 注意力替换该累加。Block AttnRes 在块级表示上做注意力,以降低大规模训练的内存与通信。随后给出跨规模的 scaling 结果,以及在 Kimi Linear(48B 总参数 / 3B 激活、1.4T Token)上的预训练与下游评测。

总结

该工作把残差连接同时看作深度方向的混合规则,而不仅是梯度通路。PreNorm 下每一层输出以系数 1 相加,残差流幅值可随深度增长,新一层更新占比下降。AttnRes 用 Softmax 权重替换单位求和;Block AttnRes 把同一机制放在块级摘要上,以便与流水线训练相容。证据包括 scaling-law 对照,以及 48B / 3B、1.4T Token 的大规模实验。这些结果支持在所报告配方下信息流得到改善,尚不能推出对所有 Dense 模型或后训练设置均成立。

详细解读

把深度混合看作序列问题

PreNorm 一步可写成 h_l = h_(l-1) + f_l(h_(l-1))。展开后,当前状态等于词嵌入加上此前所有层输出。恒等通路有利于优化,却无法在早期表示之间作出取舍。

We observe a formal duality between depth-wise accumulation and the sequential recurrence in RNNs.

这一对照是结构上的。循环网络沿时间把历史压成单一状态;残差沿深度做同样的事。序列建模用注意力替换了循环。该工作把同一替换用于深度:h_l = Σ α_(i→l) · v_i,权重由每层一个可学习伪查询 w_l 经 Softmax 得到。打分前对来源做 RMSNorm,避免仅因残差幅值较大而获得更高权重。

Indeed, standard residual connections and prior recurrence-based variants can all be shown to perform depth-wise linear attention; AttnRes generalizes them to depth-wise softmax attention, completing for depth the same linear-to-softmax transition that proved transformative over sequences.

Full AttnRes 因此不是 Token 自注意力的复制,而是对网络自身层输出的路由。Softmax 引入竞争:提高某一来源的权重,其余来源随之下降。普通残差不具备这一性质。

大规模训练为何需要 Block AttnRes

普通训练中,层输出本需保留以供反向传播,Full AttnRes 额外内存有限。一旦采用激活重计算与流水线并行,这些张量必须存储并传输,内存与通信按 O(L d) 增长。Block AttnRes 将层划分为块,块内用标准残差压缩为单一表示,注意力仅作用于块级摘要。

This brings both memory and communication down to O(N d), and together with infrastructure optimizations (§4), Block AttnRes serves as a drop-in replacement for standard residual connections with marginal training cost and negligible inference latency overhead.

块内仍为加性混合,块间对块摘要做注意力。这一设计用较粗的访问粒度换取更低的通信成本。论文报告,在跨流水线缓存与两阶段推理下,典型推理时延开销小于 2%。

实验证据

Scaling law experiments confirm that AttnRes consistently outperforms the baseline across compute budgets, with Block AttnRes matching the loss of a baseline trained with 1.25× more compute.

Scaling 实验在五种激活参数规模、8,192 Token 上下文上训练 PreNorm 基线、Full AttnRes 与 Block AttnRes。在 5.6 PFLOP/s-days 处,Block AttnRes 验证损失为 1.692,拟合基线为 1.714,论文将其解释为 1.25 倍计算优势。

Analysis of the resulting training dynamics reveals that AttnRes mitigates PreNorm dilution, with output magnitudes remaining bounded across depth and gradient norms distributing more uniformly across layers. On downstream benchmarks, our final model improves over the baseline across all evaluated tasks.

大规模实验采用 Kimi Linear:总参数 48B、激活 3B、27 个 Transformer 块,预训练与中期训练合计 1.4T Token。训练动态显示输出幅值为有界的周期模式,而非随深度单调增长,梯度范数在层间更均匀。下游任务均有提升,GPQA-Diamond、数学与代码上更为明显。

证据边界

上述结果支持在所报告架构与配方下信息流得到改善。它们尚不能推出所有 Dense 模型、优化器或后训练设置均可获得同等收益。该工作为提出方预印本,独立复现尚未见到。更低的预训练损失并不自动等于更好的指令遵循或更低的部署成本。较大的深度权重只说明该层被使用,并不构成人类可读的因果解释。Block AttnRes 是这一提议在分布式训练条件下仍然可行的形式。