滑動窗口勝過線性注意力

摘要

由於二次複雜度注意力（quadratic attention）的特性，大型語言模型（LLM）會消耗大量記憶體與能源。每增加一個新詞元，成本都比上一個更高。每個額外詞元的鍵（key）與值（value）都必須無限期地儲存在記憶體中，這是不可持續的。為了解決二次方縮放（quadratic scaling）問題，已有數種替代方案被提出，其中之一是將 LLM 改造為使用線性注意力。鑑於此方法有望以低成本解決二次方縮放問題，同時維持最先進的效能，這個想法已吸引大量關注。然而，此研究方向尚未與更簡單的基線方法進行適當比較。在這項工作中，我們證明帶有 sink 詞元的滑動視窗注意力（SWA）的效能不亞於，甚至優於後訓練的線性注意力模型。我們在多個 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.