Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Training-free Context-adaptive Attention for Efficient Long Context Modeling

Zeng You, Yaofo Chen, Shuhai Zhang, Zhijie Qiu, Tingyu Wu, Yingjian Li, Yaowei Wang, Mingkui Tan

TL;DR

The paper tackles the inefficiency of self-attention in ultra-long sequences by introducing Training-free Context-adaptive Attention (TCA-Attention), a two-phase, training-free sparse-attention mechanism. It combines an offline head-specific sparsity calibration with an online core-context selection that forms global and local token subsets, preserving important context while discarding redundancy. A formal bound on approximation error and extensive experiments show up to 2.8× speedups and 61% KV-cache reduction at 128K context, with performance comparable to full attention across benchmarks. The approach is plug-and-play, requiring no retraining or architectural changes, and is demonstrated on state-of-the-art LLMs for long-context reasoning and multi-turn dialogues.

Abstract

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.

Training-free Context-adaptive Attention for Efficient Long Context Modeling

TL;DR

The paper tackles the inefficiency of self-attention in ultra-long sequences by introducing Training-free Context-adaptive Attention (TCA-Attention), a two-phase, training-free sparse-attention mechanism. It combines an offline head-specific sparsity calibration with an online core-context selection that forms global and local token subsets, preserving important context while discarding redundancy. A formal bound on approximation error and extensive experiments show up to 2.8× speedups and 61% KV-cache reduction at 128K context, with performance comparable to full attention across benchmarks. The approach is plug-and-play, requiring no retraining or architectural changes, and is demonstrated on state-of-the-art LLMs for long-context reasoning and multi-turn dialogues.

Abstract

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 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.

Paper Structure

This paper contains 15 sections, 1 theorem, 9 equations, 6 figures, 9 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Motivations
  4. Understanding Bottlenecks of Self-Attention in Long-context Modeling
  5. Exploration and Exploitation of Attention Redundancy Properties
  6. Training-free Context-adaptive Attention
  7. Offline Head-Specific Sparsity Determination
  8. Online Core Context Selection
  9. Theoretical Guarantees for TCA-Attention
  10. Experiments
  11. Experimental Setup
  12. Comparisons with State-of-the-art Methods
  13. Ablation Studies
  14. Latency Breakdown Analysis
  15. Conclusion

Key Result

Theorem 1

Let $\gamma_i = \sum_{j \notin \mathcal{S}_{\text{total}}} \frac{\exp(\mathbf{A}_{ij})}{Z_i}$ be the total probability mass that the full softmax assigns to tokens omitted by TCA-Attention. Here, $\mathcal{S}_{\text{total}}$ denotes the complete set of token indices selected by our method, including where $|\mathbf{V}|_\infty$ is the maximum absolute entry in value matrix.

Figures (6)

  • Figure 1: Visualizations of attention distributions in LLaMA-3.1-8B-Instruct: 1) average attention scores across key positions on the first row, and 2) scatter points highlighting attention scores above row-wise averages on the second row. The observations motivate the design principles of TCA-Attention (refer to Appendix B for more visualizations).
  • Figure 2: Overview of our TCA-Attention. We dynamically select a subset of key/value tokens, which combines 1) a global subset ${\bf K}^G,{\bf V}^G$, selected online based on a pre-determined configuration (Sec. \ref{['sec:online']}), to model long‐distance dependencies; and 2) a local subset ${\bf K}^L,{\bf V}^L$, preserving neighboring tokens to capture fine‐grained local context. The concatenation of these complementary subsets is used for the final attention computation in Eq. (\ref{['eq:atten']}).
  • Figure 3: Comparisons in terms of computational and storage overhead on LLaMA3.1-8B-Instruct. Attention computation latency is the time to compute a single attention layer. "ITL" (inter token latency) is the time between generating consecutive tokens (except for the first token) in decoding.
  • Figure 4: Ablations on hyperparameters. We report the average score on LongBench-E and the decoding latency.
  • Figure 5: Performance comparison across different $\tau$ values (corresponding to different sparsity ratios) on LongBench-E tasks.
  • ...and 1 more figures

Theorems & Definitions (1)

  • Theorem 1: Error Bound for a Single Query