탄성 어텐션: 효율적인 트랜스포머를 위한 테스트 시간 적응형 희소성 비율

초록

표준 어텐션 메커니즘의 2차 복잡도는 장문맥 시나리오에서 대규모 언어 모델(LLM)의 확장성에 심각한 병목 현상을 야기합니다. 단일 모델 내에서 희소 어텐션과 전체 어텐션을 결합하는 하이브리드 어텐션 전략은 실용적인 해결책을 제공하지만, 일반적으로 정적 계산 비율(즉, 희소 대 전체 어텐션의 고정된 비율)을 사용하며 추론 시 하위 작업의 다양한 희소성 민감도에 적응하지 못합니다. 이러한 문제를 해결하기 위해 우리는 모델이 입력에 따라 전체 희소도를 동적으로 조절할 수 있는 Elastic Attention을 제안합니다. 이는 기존 사전 훈련된 모델에 경량의 Attention Router를 통합하여 각 어텐션 헤드를 다양한 계산 모드에 동적으로 할당함으로써 구현됩니다. 8xA800 GPU에서 단 12시간의 학습만으로 우리의 방법은 모델이 강력한 성능과 효율적인 추론을 동시에 달성하도록 합니다. 널리 사용되는 LLM에 대한 세 가지 장문맥 벤치마크에서의 실험을 통해 우리 방법의 우수성을 입증합니다.

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.

탄성 어텐션: 효율적인 트랜스포머를 위한 테스트 시간 적응형 희소성 비율

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

초록

Support