Attention Residuals

TL;DR

提出Attention Residuals,使用softmax注意力替代固定累积,提升模型性能。

cs.CL 🔴 高级 2026-03-16 4 次浏览
Kimi Team Guangyu Chen Yu Zhang Jianlin Su Weixin Xu Siyuan Pan Yaoyu Wang Yucheng Wang Guanduo Chen Bohong Yin Yutian Chen Junjie Yan Ming Wei Y. Zhang Fanqing Meng Chao Hong Xiaotong Xie Shaowei Liu Enzhe Lu Yunpeng Tai Yanru Chen Xin Men Haiqing Guo Y. Charles Haoyu Lu Lin Sui Jinguo Zhu Zaida Zhou Weiran He Weixiao Huang Xinran Xu Yuzhi Wang Guokun Lai Yulun Du Yuxin Wu Zhilin Yang Xinyu Zhou
深度学习 残差连接 注意力机制 模型优化 大规模训练

核心发现

方法论

本文提出了Attention Residuals (AttnRes),通过softmax注意力机制替代传统的固定累积残差连接。每层通过学习的、与输入相关的权重选择性地聚合先前层的表示。为解决大规模模型训练中内存和通信开销的问题,提出了Block AttnRes,将层划分为块,并在块级别进行注意力操作,从而减少内存占用,同时保留完整AttnRes的大部分收益。

关键结果

  • 在Kimi Linear架构中集成AttnRes后,预训练1.4T tokens,AttnRes有效缓解了PreNorm稀释问题,输出幅度和梯度分布在深度上更均匀,所有评估任务的下游性能均有提升。
  • Block AttnRes在减少内存占用的同时,性能与完整AttnRes相当,训练开销仅增加不到4%。
  • 缩放实验表明,AttnRes在不同模型规模下均表现出一致的改进,Block AttnRes在计算预算上匹配了基线的1.25倍。

研究意义

该研究通过引入Attention Residuals,解决了传统残差连接中信息累积不均的问题,显著提高了深度学习模型的训练稳定性和性能。AttnRes不仅在理论上提供了一种新的深度信息聚合方式,还在实践中展示了其在大规模模型训练中的有效性和可扩展性。

技术贡献

技术贡献包括提出了一种新的残差连接方式,通过softmax注意力机制实现深度信息的选择性聚合。Block AttnRes通过块级别的注意力操作,显著减少了内存和通信开销,提供了一种可扩展的解决方案。

新颖性

AttnRes首次将softmax注意力机制应用于深度信息的聚合,与传统的线性累积方式相比,提供了更灵活和高效的选择性聚合策略。

局限性

  • 在极大规模的模型中,尽管Block AttnRes减少了内存占用,但仍然存在一定的通信开销。
  • 在某些特定任务中,可能需要进一步优化参数以达到最佳性能。

未来方向

未来研究可以探索在不同架构中应用AttnRes的效果,以及进一步优化Block AttnRes以减少通信开销。

AI 总览摘要

在现代大型语言模型中,残差连接是标准组件,但其固定的累积方式导致了深度信息的稀释。本文提出了Attention Residuals (AttnRes),通过softmax注意力机制替代传统的固定累积方式,使每层能够选择性地聚合先前层的表示,从而缓解了信息稀释问题。

为了应对大规模模型训练中的内存和通信开销,研究引入了Block AttnRes,将层划分为块,并在块级别进行注意力操作。这种方法不仅减少了内存占用,还保留了完整AttnRes的大部分性能提升。实验表明,AttnRes在不同模型规模下均表现出一致的改进,Block AttnRes在计算预算上匹配了基线的1.25倍。

通过在Kimi Linear架构中集成AttnRes并进行大规模预训练,研究展示了其在下游任务中的显著性能提升。AttnRes有效缓解了PreNorm稀释问题,输出幅度和梯度分布在深度上更均匀。这一创新为深度学习模型的训练和性能优化提供了新的思路和方法。

深度分析

研究背景

深度学习模型中的残差连接通过允许梯度绕过变换,极大地提高了训练深度网络的稳定性。然而,传统的残差连接采用固定的累积方式,导致深度信息的稀释和隐藏状态的无控制增长。近年来,研究者们尝试通过引入可学习的输入相关权重来优化信息聚合,但这些方法在大规模模型中难以扩展。

核心问题

传统残差连接的固定累积方式导致深度信息的稀释,使得每层的贡献逐渐被削弱。随着深度的增加,隐藏状态的幅度呈O(L)增长,早期层的信息被掩盖,无法被选择性地检索。这一问题在大规模模型中尤为突出,影响了模型的训练效率和性能。

核心创新

本文提出的Attention Residuals (AttnRes)通过softmax注意力机制替代固定累积方式,使每层能够选择性地聚合先前层的表示。Block AttnRes通过块级别的注意力操作,显著减少了内存和通信开销,提供了一种可扩展的解决方案。

方法详解

  • �� 使用softmax注意力机制替代固定累积方式,实现深度信息的选择性聚合。
  • �� 引入Block AttnRes,将层划分为块,并在块级别进行注意力操作,减少内存占用。
  • �� 结合缓存管道通信和两阶段计算策略,优化大规模模型训练中的效率。

实验设计

在Kimi Linear架构中集成AttnRes,并在1.4T tokens上进行预训练。实验设计包括缩放实验、组件消融实验和下游任务评估。结果表明,AttnRes在不同模型规模下均表现出一致的改进,Block AttnRes在计算预算上匹配了基线的1.25倍。

结果分析

AttnRes有效缓解了PreNorm稀释问题,输出幅度和梯度分布在深度上更均匀。Block AttnRes在减少内存占用的同时,性能与完整AttnRes相当,训练开销仅增加不到4%。

应用场景

AttnRes可以直接应用于大型语言模型的训练中,提升模型的性能和稳定性。其块级别的注意力操作适用于需要减少内存和通信开销的大规模模型训练。

局限与展望

尽管Block AttnRes减少了内存占用,但在极大规模的模型中,仍然存在一定的通信开销。此外,在某些特定任务中,可能需要进一步优化参数以达到最佳性能。

通俗解读 非专业人士也能看懂

想象一下,传统的残差连接就像一个流水线,每个工人都在固定的时间内完成相同的工作,导致一些工人的贡献被稀释。Attention Residuals就像是一个智能的流水线管理系统,它能根据每个工人的表现来调整他们的工作时间,从而提高整体效率。Block AttnRes则是将工人分成小组,每个小组都有自己的管理系统,这样不仅减少了管理成本,还能保持高效的工作流程。

简单解释 像给14岁少年讲一样

嘿,想象一下你在学校的团队项目中,每个人都必须贡献相同的工作量,但有些人做得更好。传统的方法就像是让每个人都做相同的事情,结果有些人的努力被浪费了。Attention Residuals就像是一个聪明的队长,他能根据每个人的特长分配任务,让整个团队表现更好。Block AttnRes则是把大家分成小组,每个小组都有自己的队长,这样既省时又高效!

术语表

Attention Residuals (注意力残差)

一种替代传统残差连接的方法,通过softmax注意力机制实现深度信息的选择性聚合。

用于解决深度信息稀释问题。

Block AttnRes (块级注意力残差)

将层划分为块,并在块级别进行注意力操作,减少内存和通信开销。

用于大规模模型训练中的内存优化。

PreNorm (预归一化)

一种在深度学习模型中常用的归一化技术,通常用于提高训练稳定性。

在本文中,AttnRes缓解了PreNorm稀释问题。

Softmax Attention (软最大注意力)

一种通过softmax函数计算注意力权重的方法,使模型能够选择性地聚合信息。

用于替代传统的固定累积方式。

Kimi Linear (Kimi线性架构)

一种大规模的Mixture-of-Experts Transformer架构,用于高效的深度学习模型训练。

本文中用于验证AttnRes的有效性。

开放问题 这项研究留下的未解疑问

  • 1 如何在极大规模模型中进一步减少通信开销?
  • 2 在不同任务中,如何优化AttnRes的参数以达到最佳性能?

应用场景

近期应用

大型语言模型训练

通过引入AttnRes,提高模型的性能和稳定性,适用于需要高效训练的大型语言模型。

远期愿景

通用深度学习优化

AttnRes的引入可能会影响深度学习的整体优化策略,推动更高效的模型设计。

原文摘要

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.

cs.CL