Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models-优快云博客

本文链接：https://blog.youkuaiyun.com/c_cpp_csharp/article/details/136067645

本文介绍了Lightning Attention-2，一种新型线性注意力机制，首次实现无限序列长度处理的理论优势。通过块内和块间的线性注意力计算，解决了累积求和问题，保持了固定内存下的恒定训练速度。实验证明，该方法在任意序列长度下保持一致性能，显著优于其他机制。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

本文是LLM系列文章，针对《Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models》的翻译。

Lightning Attention-2：处理大型语言模型中无限序列长度的免费午餐

摘要
1 引言
2 相关工作
3 方法
4 实验
5 结论

摘要

线性注意力是一种有效的注意力机制，最近成为传统softmax注意力的一种很有前途的替代方法。凭借其在线性计算复杂性中处理标记的能力，线性注意力在理论上可以在不牺牲速度的情况下处理无限长度的序列，即，在固定内存消耗的情况下，保持不同序列长度的恒定训练速度。然而，由于累积求和（cumsum）的问题，当前的线性注意力算法无法在随意的环境中展示其理论优势。在本文中，我们提出了Lightning Attention-2，这是第一个使线性注意力实现其理论计算优势的线性注意力实现。为了实现这一点，我们利用了平铺的思想，在线性注意力计算中分别处理块内和块间分量。具体来说，我们对块内使用传统的注意力计算机制，并对块间应用线性注意力核技巧。通过前向和后向程序采用平铺技术，以充分利用GPU硬件。我们在Triton中实现了我们的算法，使其具有IO意识和硬件友好性。
在不同的模型大小和序列长度上进行了各种实验。Lightning Attention-2无论输入序列长度如何，都能保持一致的训练和推理速度，并且明显快于其他注意力机制。源代码位于