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

Zecheng Tang; Quantong Qiu; Yi Yang; Zhiyi Hong; Haiya Xiang; Kebin Liu; Qingqing Dang; Juntao Li; Min Zhang

Source-linked AI summary

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

Zecheng Tang, Quantong Qiu, Yi Yang, Zhiyi Hong, Haiya Xiang, Kebin Liu, Qingqing Dang, Juntao Li, Min Zhang

arXiv:2601.17367v2Jan 24, 2026cs.CLcs.AI

TL;DR

Long-context full attention is expensive, while fixed hybrid sparsity ratios do not adapt to tasks with different sparsity sensitivities. Elastic Attention adds a lightweight Attention Router that dynamically assigns heads to full or sparse attention, achieving strong benchmark performance and efficient inference across long-context settings.

Problem
Full attention has quadratic computational and memory complexity, while fixed sparse–full ratios fail to adapt to downstream tasks with different sparsity sensitivities.
Method
Elastic Attention integrates a lightweight Attention Router into a pretrained transformer and dynamically routes each attention head between full and sparse computation based on input task regimes.
Results
Across three long-context benchmarks and multiple LLMs, Elastic Attention consistently achieves strong or superior performance while enabling more efficient inference.
Takeaways & Limitations
Effective sparsity allocation can distinguish sparsity-robust from sparsity-sensitive task regimes without modifying pretrained backbone parameters.
Takeaways & Limitations
The evaluation adopts the grouped query attention definition rather than multi-head attention, reflecting the use of GQA in modern LLM architectures.

Abstract

from arXiv · show

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.

1. Introduction

Elastic Attention addresses the scalability and task-adaptation limits of long-context attention by dynamically adjusting sparsity according to input-dependent task regimes. A lightweight router assigns attention heads to full or sparse computation while preserving efficient inference.

Standard full attention has quadratic computational and memory complexity, creating a scalability bottleneck for expanding long-context sequences.
Hybrid methods combine sparse and full attention to reduce cost, but fixed proportions require extensive task-specific validation or tuning.
Downstream tasks divide into sparsity-robust tasks, such as summarization, and sparsity-sensitive tasks, such as question answering.
Elastic Attention automatically adjusts overall sparsity during prefill to accommodate these two task categories.
The Attention Router uses input hidden states for head-wise routing, assigning each head to full or sparse attention and adding 0.27M parameters per layer.
Across three LLMs, diverse long-context benchmarks, and streaming and block-sparse patterns, Elastic Attention learns task-dependent sparsity and achieves superior performance with limited training budgets.Training requires only 12 hours, although the supplied passage does not complete the associated hardware specification.

2. Preliminary

The preliminary analysis formalizes hybrid full-and-sparse attention and shows that increasing model sparsity affects downstream tasks differently. This motivates distinguishing sparsity-robust from sparsity-sensitive tasks when selecting attention computation.

Hybrid Head Mechanism: Hybrid attention assigns each head in each layer to either full attention or sparse attention, then concatenates the resulting head outputs.
Sparsity Ratio Computation: Model sparsity ratio ΩMSR measures the fraction of sparse heads, while effective sparsity also accounts for the token-pruning pattern used by those heads.
Impact of ΩMSR on Downstream Tasks: Existing hybrid mechanisms fix ΩMSR at inference time, leaving its relationship with downstream task performance insufficiently explored.
Impact of ΩMSR on Downstream Tasks: The preliminary study progressively replaces retrieval heads with sparse heads and evaluates the resulting models across six LongBench long-context tasks.
Impact of ΩMSR on Downstream Tasks: At ΩMSR = 20%, one resulting model achieves 73.85% of the backbone model’s performance, illustrating the effect of sparsity selection.
Impact of ΩMSR on Downstream Tasks: Sparsity-robust tasks remain largely insensitive to ΩMSR, whereas sparsity-sensitive tasks decline sharply after a task-dependent sparsity threshold.

3. Elastic Attention

Elastic Attention uses an input-conditioned router to assign each attention head to full or sparse computation, dynamically adapting model sparsity while preserving efficient deployment.

Elastic Attention: The Elastic Attention module integrates a lightweight Attention Router into the backbone while freezing all other model parameters during training.The router is the trainable addition to the pretrained transformer.
Attention Router: Each key–value head receives a binary full-attention or sparse-attention assignment, and the resulting head outputs are concatenated into the layer output.This head-wise routing dynamically controls the model’s overall sparsity.
Attention Router: The router uses pooled key hidden states and task-aware representations to produce head-wise logits that determine each head’s computation mode.Its task MLP infers task characteristics, while its router MLP selects full attention or sparse attention per head.
Router Optimization: Gumbel–Softmax relaxation and a straight-through estimator enable hard routing in the forward pass while preserving gradient propagation during training.Annealing sharpens the soft routing distribution toward discrete decisions.
Training Objective: The training objective combines language-modeling loss with sparsity regularization using task-dependent constraints and trainable Lagrange multipliers.The multipliers decouple sparsity–performance trade-offs across tasks and mitigate optimization conflicts.
Efficient Deployment: A fused attention kernel jointly computes retrieval and sparse heads to address runtime inefficiency from mixed head types on a single GPU.The implementation targets deployment without introducing inter-device communication overhead.

4. Experiment

Experiments evaluate Elastic Attention across real-world, extrapolation, and reasoning benchmarks with multiple LLMs and sparse-attention patterns. The method generally delivers strong or best average performance while dynamically selecting task-dependent sparsity levels.

Training Data: Training uses Qwen3-4B/8B and Llama-3.1-8B-Instruct with data spanning 8K–64K tokens and approximately 0.74B tokens total.The dataset covers sparsity-sensitive and sparsity-robust tasks, and each run completes within 12 hours on 8×A800 GPUs.
Evaluation Setting: The evaluation compares Elastic Attention with DuoAttention, PruLong, and InfLLM-V2 using Streaming Sparse Attention and XAttention patterns.Experiments use the LOOM-Eval framework and retain default hyperparameters except for XAttention’s τ = 0.9.
Real-world Long-context Tasks: On LongBench-E, Elastic Attention consistently achieves the best average performance within comparison groups while enabling more efficient inference.Its average ΩMSR is around 0.85 on code tasks and 0.68 on question-answering tasks.
Length Extrapolation Capability Testing: On RULER, the approach achieves the best performance across evaluated context lengths while maintaining ΩMSR values consistently between 0.65 and 0.7.Evaluation extends from 8K to 256K tokens despite a maximum training context length of 64K tokens.
Length Extrapolation Capability Testing: At 256K tokens, the FA–XA configuration retains strong performance because it has lower ΩESR than the other configurations.The cited comparison attributes this advantage to retaining and exploiting more relevant information.
Long-form Reasoning Task: On LongBench-v2, Elastic Attention delivers strong results across Easy and Hard settings and achieves the best average performance with either FA-SSA or FA-XA.FA–XA is not consistently strong on Qwen3-8B, indicating a model-specific configuration difference.

5. Ablation Study

The ablations examine task representation separation, head-level routing, target sparsity, and performance–efficiency trade-offs. Elastic Attention maintains strong performance while adapting sparsity across tasks and context lengths.

Attention Router: Task MLP processing significantly reduces inter-task representation similarity, indicating greater separation between task representations.Lower pairwise cosine similarity reflects improved task discrimination.
Attention Router: MLP intermediate dimensions from ×2 to ×8 achieve similar average performance, so the ×4 default provides a favorable trade-off between performance and additional parameters.The ×8 setting has the highest overall score, but ×4 is selected for the main experiments.
Attention Router: Attention routing exhibits consistent head-level specialization, with some predominantly middle-to-higher-layer heads consistently assigned to FA as retrieval heads.Other heads are frequently assigned to SA across LongBench-E tasks.
Target Sparsity: Lowering the sparsity-sensitive target from 0.7 to 0.4 increases task-level sparsity differentiation, while tsen = 0.4 can surpass the backbone but reduces inference efficiency.The main setting tsen = 0.7 and trob = 1.0 balances performance and efficiency.
Performance and Efficiency: On RULER, Elastic Attention consistently achieves the best performance, with speedup increasing at longer context lengths as it allocates higher ΩMSR.It also achieves lower ΩESR than training-based hybrid models and slightly outperforms comparable baselines such as PruLong.
Scalability: The XA–SSA setting preserves strong performance for Qwen3-4B, with an average gap within 1 point across three benchmarks, while 8B models trade degradation for substantially faster inference.XA–SSA places the entire model under a full SA regime.

6. Conclusion

Elastic Attention dynamically adjusts model sparsity from input task regimes through lightweight head-wise routing between FA and SA. Across three long-context benchmarks and multiple LLMs, it demonstrates superior performance with negligible overhead.

Conclusion: Elastic Attention dynamically adjusts model sparsity based on the input by distinguishing sparsity-robust and sparsity-sensitive task regimes.The method uses these regimes to allocate effective sparsity rather than learning separate configurations for every task.
Conclusion: A lightweight Attention Router performs head-wise routing between FA and SA modes without modifying pretrained backbones.The router introduces negligible overhead.
Conclusion: Experiments across three long-context benchmarks on cutting-edge LLMs demonstrate the superiority of Elastic Attention.The conclusion reports this outcome across the evaluated models and benchmarks.

B.2. Hybrid Efficient Architectures

Hybrid efficient architectures combine attention modes or linear-complexity operators to reduce computation while preserving capabilities. Existing approaches use increasingly fine-grained but generally static assignments of computation patterns.

Inter-layer Hybridization: Hybrid architectures integrate Full Attention with linear-complexity operators to balance efficiency and performance.Inter-layer designs interleave linear layers such as SSMs or RNNs with standard attention layers.
Inter-layer Hybridization: Large-scale inter-layer hybrids such as Jamba use fixed block-wise ratios, limiting the flexibility of computation allocation.The passage characterizes fixed ratios as part of the dominant inter-layer paradigm.
Intra-layer Hybridization: Intra-layer methods assign different attention heads to FA or Streaming Sparse Attention, while LongCat uses a static ZigZag topology with linear-complexity SSA.These approaches refine granularity but retain static architectural assignments.

C.1. Retrieval Score Calculation

The retrieval score calculation ranks attention heads by how strongly they attend to needle tokens in long-context probing. The ranking then determines which heads retain FA and which switch to SSA for a target sparsity ratio.

Retrieval Score Calculation: Retrieval Head probing inserts a key “needle” into a long “haystack” and evaluates each head’s ability to retrieve it.The procedure uses Llama-3.1-8B-Instruct and ranks heads by retrieval ability.
Retrieval Score Calculation: The retrieval score Sℓ,h measures attention mass from the last token to the needle positions, with higher scores indicating stronger activation on relevant information.The attention matrix provides weights over sequence positions.
Sparsity Assignment: Heads are ranked by retrieval score, and ΩMSR denotes the proportion converted to local attention.This ranking provides the basis for simulating different sparsity levels.
Sparsity Assignment: For target ΩMSR, k = ⌊(1 − ΩMSR) · (L · H)⌋ heads with the highest retrieval scores retain FA, while the remaining heads use SSA.The thresholding preserves global information integration through the top-ranked retrieval heads.

D.1. Training Configuration and Hyperparameters

Elastic Attention uses parameter-efficient router training and a unified block-sparse kernel to dynamically dispatch attention heads while controlling optimization overhead. The configuration combines frozen pretrained backbones, long-context training, and differentiable routing.

Model Architecture: The pretrained backbone remains frozen while only Attention Router parameters are optimized across Qwen3-4B, Qwen3-8B, and Llama-3.1-Instruct models.
Optimization Setup: Training uses 65,536-token sequences with bfloat16 precision, AdamW, FSDP hybrid sharding, and separate learning rates for router and regularization parameters.The router learning rate is 5 × 10−4, while sparsity regularization uses 1 × 10−3.
Optimization Setup: The cosine-decay schedule follows a linear warmup over the first 20% of training steps.
Evaluation: Evaluation compares Elastic Attention with training-free XAttention and training-based InfLLM v2, MoBA, NSA, PruLong, and DuoAttention under a unified benchmarking protocol.
Sparse Inference: Block Sparse Attention uses block size 64, chunk size 16,384, and 128 sink tokens for efficient streaming inference.These settings determine sparsity granularity, ultra-long-sequence processing, and streaming stability.
Sparse Inference: BSA passes routing metadata directly to one unified kernel launch, avoiding tensor splitting and reducing memory and kernel-scheduling overhead.
Router Optimization: Gumbel-Softmax, Gumbel-Sigmoid, and a straight-through estimator make discrete head selection trainable, while temperature annealing moves routing toward hard binary decisions.The decay rate is r = 0.6; high temperature supports exploration early, and low temperature supports deployment-like routing later.

G.2. Evaluation Results on LongBench-V2 and RULER

The evaluation examines LongBench-V2 and RULER across long contexts, alongside analyses of router geometry and attention-head behavior. Elastic Attention maintains stronger long-context performance and establishes a favorable accuracy–efficiency frontier through adaptive sparsity.

Benchmark Setup: RULER evaluation spans context lengths from 8K to 262K tokens, with comparative results summarized alongside LongBench-V2 experiments.
Task Geometry: Router representations for different tasks approach zero cosine similarity, indicating orthogonal local subspaces and implicit task disentanglement.
Attention Mode Routing: Qwen3 shows universally active or sparse heads across tasks, whereas Llama3.1 shifts active heads according to input context.
Adaptive Sparsity: Task-dependent sparsity levels emerge during training because Lagrangian constraints allow different tasks to tolerate different gaps from the target sparsity.
RULER Results: 68.51 is maintained by Elastic Attention FA-XA at 256K context, while XA-SSA reaches 47.68 versus XAttention’s 35.82.
Efficiency Results: XA-SSA achieves approximately 0.995 sparsity with 3.28× speedup, while FA-XA provides 1.51× acceleration with greater information retention.
Overall Comparison: Elastic Attention establishes a superior Pareto frontier, with FA-XA emphasizing retention and XA-SSA emphasizing throughput across their respective operating regimes.

H.4. Loss Curve and Monitoring metrics

Training remains stable while the router learns task-specific sparsity allocations and adaptive regularization. Router selectivity generally saturates with short boundary inputs, supporting a 100-token configuration rather than full-sequence routing.

Optimization Stability: The language-modeling loss stabilizes around 2.1, while sparsity regularization falls from approximately 0.16 to 0.06 within the first 100 steps.
Task-Dependent Allocation: Sparsity-sensitive Code and In-Context tasks converge to ΩMSR values of approximately 0.80–0.85, allocating more Full Attention to preserve performance.
Task-Dependent Allocation: Sparsity-robust Q&A plateaus near the target threshold, indicating that it can tolerate higher sparsity without unnecessary computation.
Adaptive Coefficients: The In-Context coefficient λ5 increases most aggressively, reflecting stronger pressure to satisfy density requirements for that task.
Router Input Length: Performance generally saturates around 100–200 input tokens, while full-sequence routing can produce suboptimal decisions through signal dilution.
Router Input Length: Boundary pooling uses the first and last 100 tokens to capture system instructions and user queries while filtering long-context content noise.

I. Error Analysis

Qualitative error analyses compare Elastic Attention with baselines on policy, legal, and narrative long-context tasks. The reported examples emphasize accurate retrieval of task-relevant details where baselines produce generic or hallucinated answers.

Cross-Case Interpretation: The representative outputs attribute performance improvement primarily to identifying and responding to contextual segments relevant to each query.
Policy Reasoning: On a complex policy reasoning task, Elastic Attention identifies the required Gradual approach, while baselines propose Aggressive or Immediate measures that violate the stability constraint.
Legal Understanding: On a bilingual legal document, Elastic Attention extracts the specific FAA provision on asset reallocation for public use, unlike baselines’ generic legal descriptions.
Narrative Understanding: For narrative entity tracking, Elastic Attention retrieves the correct trio of conspirators, while baselines incorrectly include Villefort, the later judicial figure.