Elastic Attention: Test-time Adaptive Sparsity Ratios for Efficient Transformers

TL;DR

Elastic Attention tackles the quadratic cost of attention in long-context transformers by introducing a test-time, data-driven sparsity mechanism. A lightweight Attention Router gates each head between full attention (FA) and sparse attention (SA), trained with a continuous relaxation via Gumbel-Softmax and a straight-through estimator, and deployed with a fused Block Sparse Attention kernel. The method distinguishes sparsity-robust from sparsity-sensitive tasks, enabling dynamic, task-aware sparsity that preserves performance while increasing inference speed, validated across LLM backbones on LongBench, LongBench-V2, and RULER with minimal backbone modifications. Ablations show the router’s design choices and sparsity targets shape performance and efficiency, revealing a favorable Pareto frontier for accuracy and throughput. Overall, Elastic Attention offers a practical, scalable path to efficient long-context reasoning in large transformers.

Abstract

The quadratic complexity of standard attention mechanisms poses a significant scalability bottleneck for large language models (LLMs) in long-context scenarios. While hybrid attention strategies that combine sparse and full attention within a single model offer a viable solution, they typically employ static computation ratios (i.e., fixed proportions of sparse versus full attention) and fail to adapt to the varying sparsity sensitivities of downstream tasks during inference. To address this issue, we propose Elastic Attention, which allows the model to dynamically adjust its overall sparsity based on the input. This is achieved by integrating a lightweight Attention Router into the existing pretrained model, which dynamically assigns each attention head to different computation modes. Within only 12 hours of training on 8xA800 GPUs, our method enables models to achieve both strong performance and efficient inference. Experiments across three long-context benchmarks on widely-used LLMs demonstrate the superiority of our method.

Figures (17)

• Figure 1: Comparison between Elastic Attention (ours) and existing approaches on LongBench-V2 bai2024longbench2. "(XA+SSA)" and "(FA+SSA)" denote our different settings.
• Figure 2: Trend of model performance as the hybrid model sparsity ratio ($\mathbf{\Omega_{MSR}}$) increases. We report model performance as a relative percentage score with respect to that of FA.
• Figure 3: Illustration of our proposed Elastic Attention. (a) shows the adapted model block with frozen backbone parameters; (b) details the dynamic assignment of heads via the Attention Router module; (c) presents the lightweight design of the Attention Router.
• Figure 4: Comparison of our fused kernel with a Torch-based sequential implementation for layer-wise hybrid attention.
• Figure 5: Visualization of task representation similarity. (Left) before Task MLP, the pooled hidden states exhibit high pairwise cosine similarity across different tasks; (Right) after passing through the Task MLP, the inter-task similarity significantly decreases.

Theorems & Definitions (2)

• Definition 2.1: Model Sparsity Ratio, $\Omega_{\mathrm{MSR}}$
• Definition 2.2: Effective Sparsity Ratio, $\Omega_{\mathrm{ESR}}$