不需训练的注意力机制,让大模型长文本处理更快更省内存。
Training-free Context-adaptive Attention for Efficient Long Context Modeling
- 通过离线校准和在线筛选,动态选择关键信息token
- 128K上下文下速度提升2.8倍,缓存减少61%仍保持性能
- 无需训练或改架构,可直接替换现有模型
大型语言模型依赖自注意力机制建模长距离依赖,但其与序列长度呈平方级增长的计算开销带来严重挑战。现有稀疏注意力和键值缓存压缩方法常受限于固定模式、无法兼顾预填充与解码阶段,或需额外训练。本文提出无需训练的上下文自适应注意力(TCA-Attention),通过两个轻量级阶段实现高效长文本推理:一为离线校准阶段,单次前向传播确定各注意力头的稀疏预算;二为在线令牌筛选阶段,利用轻量冗余度度量自适应保留核心上下文。该方法统一加速预填充与解码过程,降低键值缓存内存占用,且无需参数更新或结构改动。理论分析表明其近似误差有界。大量实验显示,在128K上下文长度下,TCA-Attention实现2.8倍加速,键值缓存减少61%,性能媲美全注意力,是高效的即插即用长文本推理方案。
原文摘要 · Abstract (English)
Large Language Models (LLMs) have demonstrated remarkable capabilities across a wide range of natural language processing tasks. These capabilities stem primarily from the self-attention mechanism, which enables modeling of long-range dependencies. However, the quadratic complexity of self-attention with respect to sequence length poses significant computational and memory challenges, especially as sequence length extends to extremes. While various sparse attention and KV cache compression methods have been proposed to improve efficiency, they often suffer from limitations such as reliance on fixed patterns, inability to handle both prefilling and decoding stages, or the requirement for additional training. In this paper, we propose Training-free Context-adaptive Attention (TCA-Attention), a training-free sparse attention mechanism that selectively attends to only the informative tokens for efficient long-context inference. Our method consists of two lightweight phases: i) an offline calibration phase that determines head-specific sparsity budgets via a single forward pass, and ii) an online token selection phase that adaptively retains core context tokens using a lightweight redundancy metric. TCA-Attention provides a unified solution that accelerates both prefilling and decoding while reducing KV cache memory footprint, without requiring parameter updates or architectural changes. Theoretical analysis shows that our approach maintains bounded approximation error. Extensive experiments demonstrate that TCA-Attention achieves a 2.8$\times$ speedup and reduces KV cache by 61% at 128K context length while maintaining performance comparable to full attention across various benchmarks, offering a practical plug-and-play solution for efficient long-context inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。