滑动窗口优于线性注意力

摘要

由于二次注意力的特性，大语言模型（LLM）会消耗大量内存和能量。每一个新词元都比前一个词元消耗更多成本。每新增一个词元，其键和值都必须无限期地存储在内存中，这一点是难以持续的。为了解决二次缩放问题，研究人员提出了多种替代方案，其中之一是将LLM改造为使用线性注意力。这一思路有望在低成本下解决二次缩放问题并保持最先进的性能，因而受到了广泛关注。然而，这一研究方向尚未与更简单的基线方法进行充分的比较。在本工作中，我们表明，带有注意力汇点的滑动窗口注意力（SWA）的表现不亚于甚至优于后训练的线性注意力模型。我们在多个LLM和各种下游任务中观察到了这一结论。在长上下文推理任务（如大海捞针测试和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.