小窗口也能保持大模型性能,关键靠新注意力机制补全信息。
RATTENTION: Towards the Minimal Sliding Window Size in Local-Global Attention Models
- 用线性注意力补充窗口外信息,解决局部注意力丢弃远距离依赖的问题。
- 512窗口大小即可媲美全注意力模型,短上下文效率提升显著。
- 适合追求高效推理的轻量级大模型部署,尤其长文本场景表现突出。
局部-全局注意力模型近年来成为标准Transformer的有力替代,有望在训练与推理效率上实现提升。然而,窗口大小的选择存在权衡:大窗口虽能保持性能但效率增益有限,小窗口则易导致性能下降。当前模型如Gemma2和Mistral采用保守设置(例如8192预训练长度中使用4096窗口)以保证性能。本文提出RATTENTION,一种融合特殊线性注意力机制的局部注意力变体,旨在捕捉窗口外的上下文信息,从而突破这一权衡边界。3B和12B规模的预训练实验表明,RATTENTION实现了更优的性能-效率平衡。其中,512窗口大小的RATTENTION在多种设置下均达到与全注意力模型相当的性能。此外,其线性注意力组件的循环特性提升了长序列处理能力,在RULER基准测试中得到验证。更重要的是,该方法未牺牲训练效率——得益于专用核实现和更小窗口,训练速度仍与当前最先进方法相当。相关Pallas内核与模型代码已开源,以促进后续研究。
原文摘要 · Abstract (English)
Local-global attention models have recently emerged as compelling alternatives to standard Transformers, promising improvements in both training and inference efficiency. However, the crucial choice of window size presents a Pareto tradeoff: larger windows maintain performance akin to full attention but offer minimal efficiency gains in short-context scenarios, while smaller windows can lead to performance degradation. Current models, such as Gemma2 and Mistral, adopt conservative window sizes (e.g., 4096 out of an 8192 pretraining length) to preserve performance. This work investigates strategies to shift this Pareto frontier, enabling local-global models to achieve efficiency gains even in short-context regimes. Our core motivation is to address the intrinsic limitation of local attention -- its complete disregard for tokens outside the defined window. We explore RATTENTION, a variant of local attention integrated with a specialized linear attention mechanism designed to capture information from these out-of-window tokens. Pretraining experiments at the 3B and 12B scales demonstrate that RATTENTION achieves a superior Pareto tradeoff between performance and efficiency. As a sweet spot, RATTENTION with a window size of just 512 consistently matches the performance of full-attention models across diverse settings. Furthermore, the recurrent nature inherent in the linear attention component of RATTENTION contributes to enhanced long-context performance, as validated on the RULER benchmark. Crucially, these improvements do not compromise training efficiency; thanks to a specialized kernel implementation and the reduced window size, RATTENTION maintains training speeds comparable to existing state-of-the-art approaches. We open-sourced our Pallas kernels along with model codes to facilitate further research effort.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。