델타 어텐션: 델타 보정을 통한 빠르고 정확한 희소 어텐션 추론

Delta Attention: Fast and Accurate Sparse Attention Inference by Delta Correction

May 16, 2025

저자: Jeffrey Willette, Heejun Lee, Sung Ju Hwang

cs.AI

초록

트랜스포머의 어텐션 메커니즘은 2차 복잡도를 가지며, 이로 인해 긴 시퀀스에 대해 높은 추론 비용과 지연 시간이 발생합니다. 그러나 어텐션 행렬은 대부분 희소(sparse)하므로, 효율적인 추론을 위해 많은 항목을 계산에서 생략할 수 있습니다. 희소 어텐션 추론 방법은 이러한 계산 부담을 줄이기 위해 고안되었지만, 성능 저하라는 문제를 동반합니다. 우리는 이러한 성능 저하의 한 가지 원인이 희소 계산이 어텐션 출력의 분포 변화(distributional shift)를 유발하기 때문임을 발견했습니다. 이 분포 변화는 디코딩 시점의 쿼리가 프리필(prefill) 단계의 적절한 키와 잘 정렬되지 못하게 하여 성능 하락을 초래합니다. 우리는 이러한 분포 변화를 교정하기 위해 간단하고 독창적이며 효과적인 절차를 제안합니다. 이 방법은 희소 어텐션 출력의 분포를 2차 어텐션에 가깝게 조정합니다. 우리의 방법은 어떤 희소 어텐션 방법에도 적용할 수 있으며, 평균 36%p의 성능 향상을 가져옵니다. 특히, 싱크 토큰(sink tokens)을 사용한 슬라이딩 윈도우 어텐션(sliding window attention) 위에 적용했을 때, 131K RULER 벤치마크에서 2차 어텐션 정확도의 88%를 회복하면서도 작은 오버헤드만 추가합니다. 또한, 우리의 방법은 전체 2차 어텐션 대비 약 98.5%의 희소성을 유지할 수 있어, 1백만 토큰 프리필을 처리할 때 Flash Attention 2보다 32배 빠른 속도를 달성합니다.

English

The attention mechanism of a transformer has a quadratic complexity, leading to high inference costs and latency for long sequences. However, attention matrices are mostly sparse, which implies that many entries may be omitted from computation for efficient inference. Sparse attention inference methods aim to reduce this computational burden; however, they also come with a troublesome performance degradation. We discover that one reason for this degradation is that the sparse calculation induces a distributional shift in the attention outputs. The distributional shift causes decoding-time queries to fail to align well with the appropriate keys from the prefill stage, leading to a drop in performance. We propose a simple, novel, and effective procedure for correcting this distributional shift, bringing the distribution of sparse attention outputs closer to that of quadratic attention. Our method can be applied on top of any sparse attention method, and results in an average 36%pt performance increase, recovering 88% of quadratic attention accuracy on the 131K RULER benchmark when applied on top of sliding window attention with sink tokens while only adding a small overhead. Our method can maintain approximately 98.5% sparsity over full quadratic attention, making our model 32 times faster than Flash Attention 2 when processing 1M token prefills.