Attention Residuals

2026年03月16日
  • 简介
    现代大语言模型(LLM)普遍采用带前置归一化(PreNorm)的残差连接,但其标准做法是将各层输出以固定的单位权重进行累加。这种均匀聚合方式会导致隐藏状态的幅值随网络深度增加而无约束地增长,进而逐步削弱每一层对最终表示的实际贡献。为此,我们提出“注意力残差”(Attention Residuals,简称 AttnRes),用基于前序各层输出的 softmax 注意力机制取代原有的固定权重累加方式,使每一层能够依据输入内容自适应地学习权重,有选择性地聚合更早层的表征。 为应对大规模模型训练中对全部前序层输出进行注意力计算所带来的显存与通信开销,我们进一步提出“分块注意力残差”(Block AttnRes):该方法将网络层划分为若干块,在块级别上构建并聚合表征,从而显著降低显存占用,同时保留绝大部分完整 AttnRes 所带来的性能增益。结合基于缓存的流水线通信机制与两阶段计算策略,Block AttnRes 成为一种实用、即插即用的标准残差连接替代方案,仅引入极小的额外开销。 缩放律(scaling law)实验表明,AttnRes 带来的性能提升在不同规模的模型上均保持一致;消融实验则验证了“依赖输入内容、沿深度方向动态选择”的机制确为关键收益来源。此外,我们将 AttnRes 集成至 Kimi Linear 架构(总参数量 480 亿,激活参数量 30 亿),并在 1.4 万亿 token 上完成预训练。结果表明,AttnRes 有效缓解了 PreNorm 下的表征稀释问题,使得模型各层输出幅值与梯度分布更加均匀,并在所有评估的下游任务上均取得性能提升。
  • 作者讲解
  • 图表
  • 解决问题
    标准残差连接(Residual Connections)与PreNorm结合虽为现代大语言模型(LLM)标配,但其固定权重的逐层累加导致隐藏状态幅值随深度指数增长,引发PreNorm稀释效应——即深层归一化失效、梯度分布不均、各层贡献失衡。该问题长期被忽视,属隐性架构缺陷,非全新问题但此前缺乏系统性建模与可学习缓解机制。
  • 关键思路
    提出Attention Residuals(AttnRes):用输入依赖的softmax注意力动态加权聚合所有前序层输出,替代固定权重求和;进一步设计Block AttnRes——将层划分为块,仅对块级压缩表征做注意力,兼顾表达力与效率;结合缓存式流水通信与两阶段计算,实现零侵入式替换标准残差。核心新意在于将残差聚合从静态、全局、等权操作,升维为动态、分层、内容感知的选择性融合。
  • 其它亮点
    1)在48B/3B稀疏激活的Kimi Linear架构上完成1.4T token预训练验证,显著改善跨层输出幅值与梯度分布均匀性;2)缩放律实验覆盖多个模型尺寸,证实增益具规模鲁棒性;3)消融实证‘内容依赖的深度选择’是关键,非简单增加参数或注意力;4)未提代码开源,但方法完全兼容现有训练框架,属工程友好型架构升级;5)值得深入:块划分策略自适应性、AttnRes与MoE路由的协同优化、其在长上下文与推理阶段的泛化表现。
  • 相关研究
    1)'On Layer Normalization in the Transformer Architecture' (Ba et al., 2016) —— PreNorm奠基工作;2)'DeepNet: Scaling Transformers to 1,000 Layers' (Huang et al., 2022) —— 通过深度缩放与调整初始化缓解深层衰减;3)'ReZero is All You Need: Fast Convergence at Large Depth' (Bai et al., 2020) —— 引入可学习标量缩放残差;4)'Stable Diffusion with Adaptive Residual Connections' (Liu et al., 2023) —— 在扩散模型中探索残差门控;5)'NormFormer: Improved Transformer Pretraining with Extra Normalization' (Shleifer et al., 2021) —— 增加中间归一化缓解PreNorm稀释。
许愿开讲
PDF
原文
点赞 收藏
向作者提问
NEW
分享到Link

提问交流

提交问题,平台邀请作者,轻松获得权威解答~

向作者提问