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

2026-01-24

AI Summary

This work addresses the limited scalability of large language models in long-context scenarios due to the quadratic complexity of standard attention mechanisms, as well as the inflexibility of existing hybrid attention approaches that employ static sparsity ratios and fail to account for varying sensitivity to sparsity across inputs and tasks. To overcome these limitations, we propose the first elastic attention mechanism capable of dynamically adjusting sparsity ratios at test time. Our approach integrates a lightweight attention router that adaptively assigns each attention head to either a sparse or dense computation path during inference. Trained for only 12 hours on 8×A800 GPUs, the model achieves both efficient inference and state-of-the-art performance across three mainstream long-context benchmarks, significantly enhancing adaptability to diverse tasks.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

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.

Problem

• attention mechanism
• sparsity
• long-context
• large language models
• test-time adaptation

Innovation

• Elastic Attention
• adaptive sparsity
• test-time adaptation
• efficient Transformers
• Attention Router