变长恒速：基于 Lightning Attention 的高效语言建模
计算与语言
2024-06-21 v2
摘要
我们提出了 Lightning Attention，这是首个在固定内存消耗下针对不同序列长度保持恒定训练速度的线性注意力实现。由于累加求和操作的问题，以往的线性注意力实现在因果设定下无法实现其理论优势。然而，通过利用不同的注意力计算策略来计算注意力的不同部分，可以有效地解决这一问题。具体而言，我们将注意力计算分为块内和块间，对块内使用常规注意力计算，对块间使用线性注意力核技巧。这消除了线性注意力计算中对 cumsum 的需求。此外，在前向和反向过程中均采用了分块技术，以充分利用 GPU 硬件。为了在保持有效性的同时提高准确性，我们引入了 TransNormerLLM (TNL)，一种为我们的 lightning attention 量身定制的新架构。我们在不同模型大小和序列长度的标准数据集和自收集数据集上进行了严格测试。TNL 比其他语言模型显著更高效。此外，基准测试结果表明，TNL 与采用传统 Transformer 结构的最先进 LLM 性能相当。源代码发布于 github.com/OpenNLPLab/TransnormerLLM。
引用
@article{arxiv.2405.17381,
title = {Various Lengths, Constant Speed: Efficient Language Modeling with Lightning Attention},
author = {Zhen Qin and Weigao Sun and Dong Li and Xuyang Shen and Weixuan Sun and Yiran Zhong},
journal= {arXiv preprint arXiv:2405.17381},
year = {2024}
}
备注
Accepted by ICML 2024. Yiran Zhong is the corresponding author. Code is released at github.com/OpenNLPLab/TransnormerLLM