We present Lightning Attention, the first linear attention implementation that maintains a constant training speed for various sequence lengths under fixed memory consumption. Due to the issue with cumulative summation operations (cumsum), previous linear attention implementations cannot achieve their theoretical advantage in a casual setting. However, this issue can be effectively solved by utilizing different attention calculation strategies to compute the different parts of attention. Specifically, we split the attention calculation into intra-blocks and inter-blocks and use conventional attention computation for intra-blocks and linear attention kernel tricks for inter-blocks. This eliminates the need for cumsum in the linear attention calculation. Furthermore, a tiling technique is adopted through both forward and backward procedures to take full advantage of the GPU hardware. To enhance accuracy while preserving efficacy, we introduce TransNormerLLM (TNL), a new architecture that is tailored to our lightning attention. We conduct rigorous testing on standard and self-collected datasets with varying model sizes and sequence lengths. TNL is notably more efficient than other language models. In addition, benchmark results indicate that TNL performs on par with state-of-the-art LLMs utilizing conventional transformer structures. The source code is released at github.com/OpenNLPLab/TransnormerLLM.

翻译：本文提出闪电注意力（Lightning Attention），这是首个在固定内存消耗下对多种序列长度保持恒定训练速度的线性注意力实现。由于累积求和运算（cumsum）存在的问题，以往的线性注意力实现无法在因果设定下达到其理论优势。然而，通过采用不同的注意力计算策略来处理注意力的不同部分，该问题可以得到有效解决。具体而言，我们将注意力计算划分为块内计算与块间计算：对块内计算采用传统注意力计算方法，对块间计算则采用线性注意力核技巧。这一设计消除了线性注意力计算中对cumsum运算的需求。此外，我们在前向与反向传播过程中均采用分块技术，以充分发挥GPU硬件的计算潜力。为在保持高效性的同时提升模型精度，我们进一步提出了TransNormerLLM（TNL）这一专门适配闪电注意力机制的新型架构。我们在标准数据集及自收集数据集上，针对不同模型规模与序列长度进行了严格测试。实验表明，TNL相比其他语言模型具有显著更高的效率。同时，基准测试结果显示，TNL的性能与采用传统Transformer结构的先进大语言模型相当。相关源代码已发布于github.com/OpenNLPLab/TransnormerLLM。