Various Lengths, Constant Speed: Efficient Language Modeling with Lightning Attention

我们提出了闪电注意力（Lightning Attention），这是第一个在固定内存消耗下，针对不同序列长度保持恒定训练速度的线性注意力实现。由于累积求和操作（cumsum）存在问题，以前的线性注意力实现在非正式设置下无法发挥其理论优势。然而，通过利用不同的注意力计算策略来计算注意力的不同部分，可以有效解决这个问题。具体而言，我们将注意力计算分为块内和块间两部分，并使用传统的注意力计算方法来计算块内的部分，使用线性注意力核技巧来计算块间的部分。这消除了线性注意力计算中对cumsum的需求。此外，我们采用平铺技术来充分利用GPU硬件，通过正向和反向过程。为了提高准确性并保持有效性，我们引入了TransNormerLLM（TNL），这是一种专门为我们的闪电注意力量身定制的新架构。我们对标准和自采集数据集进行了严格测试，包括不同模型大小和序列长度。TNL比其他语言模型更加高效。此外，基准测试结果表明，使用传统的Transformer结构的最先进的LLMs和TNL的性能相当。源代码已发布在github.com/OpenNLPLab/TransnormerLLM。

该论文旨在解决线性注意力机制在长序列下的累加操作问题，提出了一种新的线性注意力实现方法，即 Lightning Attention。

论文中的关键思路是将注意力计算分为 intra-blocks 和 inter-blocks 两部分，分别使用传统的注意力计算和线性注意力核技巧，从而消除了线性注意力计算中的累加操作。同时，采用了瓦片技术来充分利用 GPU 硬件。

论文采用了 TransNormerLLM（TNL）架构来提高模型的准确性和效率，并在标准和自收集的数据集上进行了严格测试。实验结果表明，TNL 比其他语言模型更加高效，并且在性能上与使用传统 Transformer 结构的最先进的 LLMs 相当。论文开源了代码。

最近相关研究包括：《Attention Is All You Need》、《BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding》等。