Elastic Attention: 효율적인 트랜스포머를 위한 테스트 시점 적응형 희소성 비율

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

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

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.