动态分配注意力预算,让关键词获取全局信息,提升长文本处理效率。
LoGo: Token-Level Dynamic Local-Global Attention
- 按每个词的需求动态分配局部与全局注意力范围
- 在相同计算开销下,长程任务性能优于全注意力模型
- 适合需要高效处理长文本的场景,如文档理解
随着上下文长度增长,注意力机制成为大语言模型的主要计算瓶颈。标准Transformer虽强大但计算效率低,因对所有词分配相同的注意力预算,忽略其上下文需求差异。现有局部-全局混合方法虽更高效,但通常在层或头间静态分配注意力范围。为此,我们提出LoGo,一种基于词粒度的动态局部-全局注意力机制,将注意力跨度直接作为预算分配的代理。每个LoGo层包含耦合的局部和全局分支:所有词均接受受限窗口内的高效局部注意力,而仅对需要长距离信息的词通过学习门控激活全上下文全局注意力。基于阈值的预算控制器在无辅助损失下维持目标全局比例,渐进掩码调度则在稀疏路由生效前稳定训练。我们进一步实现了查询稀疏的Triton内核,将减少的全局注意力计算转化为实际加速。大量实验验证了LoGo的有效性,其在不同模型规模下保持了全注意力Transformer的扩展特性。在受控对比中,LoGo优于全注意力Transformer和同预算静态混合模型,尤其在长程检索任务中表现显著。分析表明,LoGo可学习可解释的跨度分配模式。结果表明,学习得到的词级跨度分配是提升长上下文性能-计算权衡的有效且可扩展的方法。
原文摘要 · Abstract (English)
As context lengths scale, attention increasingly becomes a primary computational bottleneck in large language models. Standard Transformers remain powerful but computationally inefficient, as they allocate the same attention budget to every token regardless of its contextual demand. Existing local-global hybrids provide a more efficient alternative by mixing restricted- and full-context attention, but they typically allocate span statically across layers or heads. To address these limitations, we propose LoGo, a token-level dynamic local-global attention mechanism that uses attention span as a direct proxy for attention budget allocation. Each LoGo layer contains coupled local and global branches: all tokens receive efficient local attention over a restricted context window, while a learned gate activates global attention with full-context access only for tokens requiring long-range information. A threshold-based budget controller maintains a target global ratio without auxiliary losses, and a progressive masking schedule stabilizes training before sparse routing takes effect. We further implement query-sparse Triton kernels that convert reduced global-attention computation into practical speedups. Extensive experiments validate LoGo's effectiveness, showing that it preserves the scaling behavior of full-attention Transformers across model sizes. In controlled comparisons, LoGo improves over the full-attention Transformer and matched-budget static local-global hybrids, with clear gains on long-range retrieval. Analysis further shows that LoGo learns interpretable span allocation patterns. These results suggest that learned token-level span allocation is an effective and scalable way to improve the long-context performance-compute trade-off.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。