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

本文是LLM系列文章，针对《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无论输入序列长度如何，都能保持一致的训练和推理速度，并且明显快于其他注意力机制。源代码位于Lightning Attention-2。

1 引言

2 相关工作

3 方法

4 实验

5 结论

在本文中，我们介绍了Lightning Attention-2，这是线性注意力的一种开创性实现，它有效地利用了其理论计算优势，特别是在因果环境中。我们的方法采用了“分而治之”和平铺技术的概念，成功地解决了当前线性注意力算法的局限性，特别是与累积求和相关的挑战。通过将计算分为块内和块间组件，我们可以有效地充分利用GPU硬件的潜力，确保效率。我们在各种模型大小和序列长度上的广泛实验表明，Lightning Attention-2不仅在不考虑输入序列长度的情况下保持了一致的训练速度，而且在速度和准确性方面优于现有技术的注意力机制。这一突破对大型语言模型的未来有着深远的影响，尤其是那些需要处理长序列的模型。展望未来，我们打算将序列并行性与Lightning Attention-2结合使用，旨在促进超长序列的训练，有效克服现有的硬件限制。