让大模型学会何时用长程注意力，显著提升长文本处理效率
Learning When to Attend: Conditional Memory Access for Long-Context LLMs
- 按需触发全局注意力，仅对必要token使用长距离依赖
- 将上下文长度从32K扩展至128K，性能损失小于3%
- 支持训练加速和内存压缩，适合部署长文本任务
语言模型难以泛化到预训练之外的上下文长度，限制了长时序推理与检索能力。持续在长文本数据上预训练虽有帮助，但因注意力机制的二次方复杂度而代价高昂。我们观察到，多数词元无需在整个序列上进行全局注意力，可依赖局部上下文。基于此提出L2A（Learning To Attend），通过逐词判断是否启用全局注意力，实现条件性长程记忆访问。在Qwen 2.5与Qwen 3模型上评估，有效上下文长度从32K扩展至128K，性能相较标准长上下文训练仅低3%，同时跳过约80%词元的全局注意力。我们设计专用Triton内核，在GPU上高效实现该机制，训练吞吐量与首字生成时间相比FlashAttention提升约2倍。此外，L2A支持后训练中对稀疏全局注意力层剪枝，使KV缓存内存减少高达50%且性能损失极小。代码已开源于https://github.com/awslabs/hybrid-model-factory/tree/main/examples/research/L2A。
原文摘要 · Abstract (English)
Language models struggle to generalize beyond pretraining context lengths, limiting long-horizon reasoning and retrieval. Continued pretraining on long-context data can help but is expensive due to the quadratic scaling of Attention. We observe that most tokens do not require (Global) Attention over the entire sequence and can rely on local context. Based on this, we propose L2A (Learning To Attend), a layer that enables conditional (token-wise) long-range memory access by deciding when to invoke global attention. We evaluate L2A on Qwen 2.5 and Qwen 3 models, extending their effective context length from 32K to 128K tokens. L2A matches the performance of standard long-context training to within 3\% while skipping Global Attention for $\sim$80\% of tokens, outperforming prior baselines. We also design custom Triton kernels to efficiently implement this token-wise conditional Attention on GPUs, achieving up to $\sim$2$\times$ improvements in training throughput and time-to-first-token over FlashAttention. Moreover, L2A enables post-training pruning of highly sparse Global Attention layers, reducing KV cache memory by up to 50\% with negligible performance loss. Our code is released under Apache 2.0 at https://github.com/awslabs/hybrid-model-factory/tree/main/examples/research/L2A.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。