Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models

线性注意力是一种高效的注意力机制，最近已经成为传统 softmax 注意力的一种有前途的替代方案。线性注意力的能力在于以线性的计算复杂度处理标记，理论上可以处理长度无限的序列而不会牺牲速度，即在固定的内存消耗下，对于各种序列长度维持恒定的训练速度。然而，由于累积求和（cumsum）的问题，当前的线性注意力算法无法在因果设置中展示其理论优势。本文提出了 Lightning Attention-2，这是第一个能够实现线性注意力理论计算优势的线性注意力实现。为了实现这一点，我们利用平铺的思想，分别处理线性注意力计算中的内部块和外部块组成部分。具体来说，我们利用传统的注意力计算机制来处理内部块，并应用线性注意力核技巧来处理外部块。在前向和反向过程中采用平铺技术，以充分利用 GPU 硬件。我们在 Triton 中实现了我们的算法，使其具有 IO 感知和硬件友好性。在不同的模型大小和序列长度上进行了各种实验。Lightning Attention-2 不管输入序列长度如何，都保持着一致的训练和推理速度，并且比其他注意力机制快得多。源代码可在 https://github.com/OpenNLPLab/lightning-attention 上获得。

本论文试图解决linear attention在因果场景下的累计求和问题，以实现理论上的计算优势。

通过使用分块和瓦片技术，将linear attention的计算分为内部块和外部块两部分，分别采用传统的attention计算和linear attention的计算方法，以实现在因果场景下的高效计算。

论文提出了Lightning Attention-2，是第一个能够在因果场景下实现linear attention理论计算优势的实现方法，具有IO感知和硬件友好的特点。实验结果表明，Lightning Attention-2在处理不同模型大小和序列长度时具有一致的训练和推理速度，并且比其他注意力机制更快。论文代码开源。

最近的相关研究包括《Transformers without Tears: Improving the Normalization of Self-Attention》和《Rethinking Attention with Performers》等。