让Transformer在推理时动态调整注意力稀疏度,提升长文本处理效率。
Elastic Attention: Test-time Adaptive Sparsity Ratios for Efficient Transformers
- 引入轻量级路由模块,按输入动态分配注意力头的计算模式。
- 仅用12小时训练即可在8张A800上实现高效推理与强性能。
- 适合需要长上下文推理且资源受限的场景,如手机端大模型应用。
标准注意力机制的二次复杂度成为大语言模型在长上下文场景下的显著扩展瓶颈。虽然混合注意力策略通过在单个模型中结合稀疏与全连接注意力提供了可行方案,但通常采用静态计算比例(即稀疏与全连接注意力的固定比例),无法根据下游任务在推理时对稀疏性的不同敏感性进行自适应调整。为解决此问题,我们提出Elastic Attention,使模型能够根据输入动态调整整体稀疏度。该方法通过在现有预训练模型中集成一个轻量级注意力路由模块,动态分配每个注意力头至不同计算模式。仅需在8张A800 GPU上训练12小时,即可实现优异性能与高效推理。在三个广泛使用的长上下文基准测试中,所提方法均表现出优越性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。