スライディングウィンドウが線形アテンションに勝る

要旨

二次アテンションの性質上、大規模言語モデル（LLM）は膨大なメモリとエネルギーを消費する。新しいトークンは常に前のトークンよりもコストがかかる。トークンが追加されるたびに、そのキーとバリューをメモリに無期限に保存する必要があり、この仕組みは持続不可能である。この二次スケーリング問題を解決するため、いくつかの代替手法が提案されてきた。その一つが、既存のLLMを線形アテンション用に後付けで改造する方法である。このアイデアは、低コストで最先端の性能を維持しつつ二次スケーリング問題を解決できる可能性から、大きな注目を集めている。しかし、この研究方向は、より単純なベースラインと適切に比較されていない。本研究では、アテンションシンクを備えたスライディングウィンドウアテンション（SWA）が、後訓練された線形アテンションモデルと同等以上の性能を発揮することを示す。この結果は、複数のLLMとさまざまな下流タスクにわたって観察された。長文脈推論タスク（Needle-in-a-HaystackやBABILong）では、SWAは線形アテンションよりもはるかに高い性能（2～10倍）を達成する。SWAは後訓練が不要で、極めて高速であり、必要なメモリも少ない。したがって、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.