슬라이딩 윈도우는 선형 어텐션을 능가한다

초록

이차 주의 메커니즘(quadratic attention)의 특성상, 대규모 언어 모델(LLM)은 많은 메모리와 에너지를 소비한다. 새로운 토큰이 추가될 때마다 이전 토큰보다 더 많은 비용이 발생하며, 토큰이 추가될 때마다 키(key)와 값(value)은 무기한 메모리에 저장되어야 하므로 지속 가능하지 않다. 이러한 이차 확장(quadratic scaling) 문제를 해결하기 위해 여러 대안이 제안되었으며, 그중 하나는 LLM에 선형 주의(Linear Attention)를 적용하는 것이다. 이 아이디어는 낮은 비용으로 최첨단 성능을 유지하면서 이차 확장 문제를 해결할 수 있다는 가능성 때문에 많은 주목을 받아왔다. 그러나 이 연구 방향은 더 단순한 기준 모델(baseline)과 제대로 비교된 적이 없다. 본 연구에서는 싱크(sink)를 포함한 슬라이딩 윈도우 주의(Sliding Window Attention, SWA)가 사후 학습된(post-trained) 선형 주의 모델과 동등하거나 더 나은 성능을 보인다는 것을 입증한다. 우리는 다양한 다운스트림 태스크에서 여러 LLM을 대상으로 이러한 결과를 확인했다. 장기 문맥 추론 태스크(Needle-in-a-Haystack 및 BABILong)에서 SWA는 선형 주의보다 압도적으로 높은 성능(2배에서 10배)을 달성했다. SWA는 사후 학습이 필요 없고, 매우 빠르며, 낮은 메모리를 요구하므로, 매우 저렴하고 신뢰할 수 있는 솔루션이다. 추론 메모리 비용을 줄이기 위해, 우리는 선형 모델을 사후 학습하는 대신 SWA로 전환할 것을 강력히 권장한다. 선형 주의 모델이 어느 정도 가능성을 보여주었지만, SWA와 필적하려면 처음부터 학습하거나 광범위한 사후 학습이 필요할 것으로 보인다.

English

Due to the nature of quadratic attention, Large Language Models (LLMs) consume a lot of memory and energy. Every new token costs more than the previous one. For each additional token, the keys and values must be stored in memory indefinitely, which is unsustainable. Several alternatives have been proposed to fix the quadratic scaling problem, one of which is retrofitting LLMs to use Linear Attention. This idea has attracted a lot of attention, given its promise to solve the quadratic scaling problem with state-of-the-art performance at low cost. However, this line of research has not been properly compared to simpler baselines. In this work, we show that Sliding Window Attention (SWA) with sinks performs as well or better than post-trained Linear Attention models. We observe this across multiple LLMs on various downstream tasks. For long-context reasoning tasks (Needle-in-a-Haystack and BABILong), SWA achieves massively higher performance (2 to 10 times higher than linear attention). SWA requires no post-training, is extremely fast, and requires low memory; therefore, making it an extremely cheap and reliable solution. To reduce inference memory cost, we strongly recommend switching to SWA instead of post-training linear models. Linear attention models may have shown some promise, but they likely require to be trained from scratch or extensive post-training in order to even match SWA.