让每层网络自适应选择之前层的输出,提升模型深度表现
Attention Residuals
- 用软注意力机制替代固定权重叠加,实现动态选择历史层输出
- 在480亿参数模型上训练1.4万亿词,下游任务性能全面提升
- 适合追求深层模型稳定性和性能优化的研究者与工程师
现代大语言模型中,预归一化残差连接虽为标准配置,但其对所有层输出采用固定单位权重叠加,导致隐藏状态随深度增长失控,逐步稀释各层贡献。本文提出注意力残差(AttnRes),以软注意力机制取代固定叠加,使每一层能基于输入内容,自适应地选择性聚合早期表示。为降低大规模训练时全序列注意力带来的内存与通信开销,引入块级注意力残差(Block AttnRes),将层分块后仅在块级别进行注意力计算,显著减少内存占用,同时保留大部分性能增益。结合基于缓存的流水线通信与两阶段计算策略,该方法可作为标准残差连接的零成本替换。缩放定律实验表明,性能提升在不同模型规模下一致;消融实验证明内容依赖的深度选择机制有效。进一步将AttnRes集成至Kimi Linear架构(总参数480亿,激活参数30亿),在1.4万亿词数据上预训练,成功缓解预归一化稀释问题,使各层输出幅度与梯度分布更均匀,所有下游任务表现均获提升。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。