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》的翻译。

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

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

摘要

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