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 上获得。

本论文旨在解决线性注意力机制在因果设置下的累积求和问题，从而实现其理论上的计算优势。

该论文提出了一种名为Lightning Attention-2的线性注意力实现方法，通过利用平铺思想，分别处理线性注意力计算中的内部块和外部块组件，从而实现线性注意力的理论计算优势。

该论文的实验结果表明，Lightning Attention-2无论输入序列长度如何，都能保持一致的训练和推理速度，并且比其他注意力机制更快。该论文的代码已经开源。

最近的相关研究包括《Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context》、《Sparse Transformer: Generating Long Sequences with Sparse Transformers》等。