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

Abstract: The quadratic complexity of standard attention mechanisms poses a significant scalability bottleneck for 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.

• The paper introduces Elastic Attention, a framework that dynamically allocates sparsity across transformer heads based on input and task sensitivity.
• The methodology employs a lightweight Attention Router using Gumbel-Softmax and a straight-through estimator to switch between full and sparse attention modes efficiently.
• Empirical results show that adaptive sparsity preserves near-baseline performance on robust tasks while mitigating degradation on sparsity-sensitive tasks in long-context benchmarks.

Elastic Attention: Test-time Adaptive Sparsity for Efficient Transformers

Motivation and Problem Statement

The scalability bottleneck of standard dot-product attention ($O(n^2)$ memory and compute) underlies the difficulty of deploying LLMs for long-context reasoning and real-world scenarios. While full attention (FA) preserves associative recall for factuality and reasoning, sparse attention (SA)—including block-sparse and streaming sparse patterns—enables linear or sub-quadratic computational scaling. Hybrid attention schemes that statically combine FA and SA are widely adopted, yet fixed sparsity ratios fail to accommodate the spectrum of downstream task sensitivities: some tasks (summarization, code) remain robust against aggressive sparsification, while others (multihop QA, retrieval) suffer performance collapse if FA is underprovisioned.

Elastic Attention addresses this core limitation by enabling dynamic, test-time adaptation of head-wise sparsity—realized via a lightweight, input-conditioned Attention Router—thus allowing transformer architectures to allocate sparsity in an input- and task-adaptive manner. This approach decouples the sparsity allocation policy from model pre-training, providing fine-grained automatic trade-offs between accuracy and efficiency.

Architectural Design and Mechanism

Elastic Attention integrates a head-level router module into the Transformer block (Figure 1). The router operates on pooled key hidden states and consists of two MLPs: Task MLP, which disentangles sequence-wise structure into task-relevant representations, and Router MLP, which produces per-head logits determining routing to either FA (retrieval head) or SA (sparse head).

Hard binary routing decisions are implemented during both training and inference, enabled by a Gumbel-Softmax continuous relaxation and a straight-through estimator (STE) for gradient flow. During a forward pass, heads within the same layer may be dynamically routed to distinct modes depending on the input. Crucially, the router parameter count per layer is minimal ($<$0.3M), supporting efficient inference even on commodity hardware.

Figure 1: Architectural overview of Elastic Attention showing parameter freezing, dynamic routing, and the router's lightweight module design.

To maximize kernel efficiency, a fused attention kernel processes FA and SA heads in parallel, bypassing the serial dispatch and memory overhead inherent in naive PyTorch/Torch attention modules (Figure 2).

Figure 2: Fused kernel achieves superior speedup over sequential hybrid attention implementations, enabling efficient joint computation.

Empirical Characterization of Sparsity Sensitivity

An extensive empirical study reveals sharp differences across task categories as the model sparsity ratio ($\Omega_\mathrm{MSR}$) increases (Figure 3):

• Sparsity-robust tasks (summarization, code completion) maintain near-baseline performance even at high $\Omega_\mathrm{MSR}$.
• Sparsity-sensitive tasks (single/multidoc QA, multihop reasoning) exhibit abrupt performance collapse beyond $\Omega_\mathrm{MSR} \approx 0.2-0.3$.
• All downstream tasks can be mapped into these two broad regimes, allowing a single router to allocate sparsity accordingly rather than requiring manual, task-specific tuning.

  Figure 3: Task performance sharply degrades in sparsity-sensitive tasks as $\Omega_\mathrm{MSR}$ increases, motivating adaptive allocation.

Optimization and Training Protocol

With backbone parameters strictly frozen, only router weights are updated. The training objective is a composite minimax function balancing standard LM cross-entropy loss and sparsity regularization. Soft per-task sparsity targets $\boldsymbol{t}$ specify lower and upper bounds—a loose, non-tight constraint used to improve generalization.

Router head assignments are supervised by the Gumbel-Softmax relaxation, with annealing schedules and gradient propagation facilitated by STE. Lagrangian multipliers, learned alongside router weights, decouple sparsity-performance trade-offs for each task. Training is conducted on Qwen3-4B, Qwen3-8B, and Llama-3.1-8B-Instruct backbones; full runs complete within 12 hours on 8 $\times$ A800 GPUs ($\sim$0.74B tokens).

Numerical Results and Performance Analysis

Elastic Attention was benchmarked on three core long-context suites: LongBench-E, RULER, and LongBench-V2. The results across all settings consistently favor Elastic Attention over comparable static and training-adaptive sparsity methods. On Qwen3-series and Llama-3.1-8B-Instruct:

• LongBench-E: Achieves best or near-best average scores, with dynamic allocation of $\Omega_\mathrm{MSR}$ per task (e.g., $\sim$0.85 for code, $\sim$0.68 for QA).
• Length Extrapolation (RULER 8K-256K): FA-XA configuration yields minimal degradation even at extreme sequence lengths (64K-256K tokens), outperforming static and block-attention baselines.
• LongBench-V2: Consistently achieves superior or leading scores across Easy and Hard settings against all methods, including InfLLM-V2, DuoAttention, PruLong, NSA, MoBA.

  Figure 4: Elastic Attention demonstrates superior accuracy-efficiency tradeoff on LongBench-V2 across multiple backbone architectures.

  

Figure 5: Robust length extrapolation capability and efficient sparsification at up to 256K tokens, demonstrating high accuracy and speedup.

Ablations and Theoretical Insights

The core mechanism—disentangled task representations driving adaptive sparsity—was verified by latent geometry analysis (Figure 6):

• Task MLP significantly reduces inter-task cosine similarity in the pooled latent space, separating tasks for optimal router decisions.

  Figure 6: Task MLP substantially increases task separability in hidden representations, enabling discriminative head routing.

Head-level activation heatmaps confirm that retrieval heads cluster in higher transformer layers and are robustly assigned FA; sparse heads are routed to SA (Figure 7).

Figure 7: Routing heatmap visualizes consistent FA assignment for retrieval heads across multiple tasks, with other heads routed to sparse attention.

Performance and efficiency trade-off curves (Figure 8, Figure 5) further demonstrate that Elastic Attention establishes a superior Pareto frontier in accuracy vs. sparsity and speedup.

Implications and Future Directions

Elastic Attention introduces a generalizable framework for test-time efficient LLM inference in arbitrarily long contexts:

• Practical Impact: Substantial cost reduction for long-context inference with negligible sacrifice of quality on sparsity-robust tasks. Dynamic head-wise allocation optimizes resource usage, making LLM deployment economically viable in production scenarios.
• Theoretical Perspective: Provides a principled method for input- and task-adaptive context processing, opening new lines of research in transformer capacity partitioning, modular routing, and sparsity-robust optimization.
• Future Extensions: The attention router's architecture could be further augmented with meta-learning or prompt-conditioned adaptation. Exploring extensions to multi-modal input, retrieval-augmented models, and reinforcement learning for sparsity-policy optimization are promising directions.

Conclusion

Elastic Attention achieves efficient, input-adaptive sparsification in frozen LLMs by dynamically routing attention heads based on task-sensitive representations. The method delivers strong accuracy-efficiency parities across benchmarks, with minimal parameter footprint and compatible with existing inference stacks. These findings establish Elastic Attention as a robust and general solution for scalable transformer inference in diverse, long-context applications.