Lightning Attention-2 新一代注意力机制

大语言模型在处理超长序列时，一直有个绕不开的坎儿——计算复杂度会随着序列长度二次增长，这就像你往杯子里倒水，杯子越大，倒满的时间就成倍增加。换句话说，当序列长度从1K增加到16K，需要的计算资源可不是16倍，而是256倍。这让多轮对话、长文本理解、多模态数据处理等应用都受到了严重制约。那么，有没有办法让这个复杂度变成线性增长呢？

过去的研究大多集中在如何让模型在推理阶段“适应”更长的序列。例如采用Alibi或类似相对位置编码让模型自适应不同输入长度，或者对RoPE等相对位置编码进行差值，在已有模型上再做短暂精调来扩增序列。但这些方法只是让大模型“有”了长序列处理能力，实际的训练和推理开销并没有减少——该花的钱还是得花，该等的时间还是得等。

OpenNLPLab团队试图一劳永逸地解决这个问题。他们提出并开源了Lightning Attention-2，一种新型线性注意力机制。核心效果是：在遇到显存瓶颈之前，无论序列多长，训练和推理成本都和1K序列长度时差不多。这意味着，你可以无限增大序列长度，但训练速度几乎不受影响。同时，超长文本的推理成本也与1K tokens一致，甚至更少。换句话说，大语言模型的长序列训练和推理，不再是一个“烧钱”的事儿了。如图1所示，在400M、1B、3B模型规模下，FlashAttention2加持的LLaMA训练速度随序列长度增加快速下降，而Lightning Attention-2加持的TransNormerLLM的速度几乎纹丝不动。

图1

Lightning Attention-2被知名AI博主AK转发，并入选Hugging Face每日必读论文Daily Papers之一。

Lightning Attention-2简介

让大模型的预训练速度在不同序列长度下保持一致，听起来像是不可能完成的任务。但如果注意力机制的计算复杂度相对于序列长度保持线性关系，那么这件事就能实现。自2020年线性注意力提出以来，研究人员一直在努力让它的实际效率符合理论线性复杂度。在2023年之前，大部分工作集中在把精度对齐到Transformer上。到了2023年中期，改进的线性注意力在精度上终于可以媲美最先进的Transformer架构了。然而，线性注意力中那个将复杂度变为线性的关键“左乘变右乘”trick（如图2所示），在实际实现中却远慢于直接左乘。原因在于右乘需要用到包含大量循环操作的累积求和（cumsum），大量的IO操作让效率大打折扣。

图2

为了更好地理解Lightning Attention-2的思路，先回顾一下传统softmax attention的计算公式：O=softmax((QK^T)⊙M)V，其中Q、K、V、M、O分别是query、key、value、mask和输出矩阵。mask在单向任务（如GPT）中是一个下三角全1矩阵，在双向任务（如Bert）中可以忽略。我们把Lightning Attention-2的整体思路总结为三点：

1. 线性注意力的核心之一是去掉成本高昂的softmax算子，使公式变成O=((QK^T)⊙M)V。但单向任务中mask矩阵M的存在，使得该形式依然只能左乘，无法获得O(N)复杂度。而双向任务没有mask，公式可以进一步简化为O=(QK^T)V。利用矩阵乘法结合律，可以转化为O=Q(K^T V)，我们称之为右乘（左乘对应前者）。如图2所示，线性注意力在双向任务中可以轻松达到诱人的O(N)复杂度。

2. 但decoder-only GPT形式逐渐成为LLM事实标准后，如何利用线性注意力的右乘特性加速单向任务便成了难题。为此我们提出了“分而治之”的思想：将注意力矩阵的计算分为对角阵和非对角阵两种形式，分别采用不同方式计算。如图3所示，Lightning Attention-2利用计算机领域常用的Tiling思想，将Q、K、V矩阵分别切分成相同数量的块（blocks）。块内（intra-block）计算由于mask矩阵的存在，依然保留左乘，复杂度O(N^2)；块间（inter-block）计算没有mask，可以采用右乘，享受O(N)复杂度。两者计算完成后直接相加得到第i块的输出Oi，同时对KV状态进行累加以供下一块使用。整个算法复杂度是intra-block O(N^2)与inter-block O(N)的trade-off，而block size决定了如何取得更好的平衡。

3. 细心的读者会发现，以上只是算法部分。之所以取名“Lightning”，是因为充分考虑在GPU硬件上的执行效率。受FlashAttention系列启发，实际计算时将切分后的Qi、Ki、Vi从速度慢容量大的HBM搬运到速度快容量小的SRAM上计算，大幅减少memory IO开销。该块计算完成后，输出Oi再搬回HBM。重复这一过程直到所有块处理完毕即可。

更多细节可参考论文中的Algorithm 1和Algorithm 2，它们区分了前向和反向过程，帮助深入理解。

图3

Lightning Attention-2精度对比

研究人员首先在400M参数模型上对比了Lightning Attention-2与Lightning Attention-1的精度，二者几乎没有差别。随后在1B、3B规模上，将Lightning Attention-2加持的TransNormerLLM（TNL-LA2）与其它非Transformer架构及FlashAttention2加持的LLaMA在相同语料下对比。TNL-LA2与LLaMA保持了相似趋势，且loss表现更优。这说明Lightning Attention-2在语言建模精度上完全不输最先进的Transformer架构。

在大语言模型常见任务上，研究还对比了TNL-LA2 15B与类似大小的Pythia模型。在吃掉了相同tokens的条件下，TNL-LA2在常识推理和多项选择综合能力上均略高于基于softmax注意力的Pythia模型。

Lightning Attention-2速度对比

研究人员对Lightning Attention-2与FlashAttention2进行了单模块速度和显存占用对比。相比Lightning Attention-1和FlashAttention2，Lightning Attention-2的速度表现出严格的线性增长。显存占用方面，三者趋势类似，但Lightning Attention-2占用更小，原因是FlashAttention2和Lightning Attention-1的显存占用其实也是近似线性的。

值得注意的是，这篇文章主要关注训练速度，实现了任意长序列与1K序列相似的训练速度。推理速度方面介绍不多，因为线性注意力在推理时可以无损转化为RNN模式，单token推理速度恒定。对于Transformer，当前token的推理速度与之前token数量相关。我们测试了Lightning Attention-1加持的TransNormerLLM-7B与其他7B模型的推理速度对比，吞吐量是百川的4倍，ChatGLM的3.5倍以上，显示出了优异的推理速度优势。

TransNormerLLM 15B（集成Lightning Attention-2）的最新Benchmark结果

小结

Lightning Attention-2代表了线性注意力机制的重大进步，在精度和速度上都能完美替换传统softmax注意力。它为越来越大的模型提供了可持续扩展的能力，开辟了一条更高效率处理无限长序列的途径。OpenNLPLab团队未来将研究基于线性注意力的序列并行算法，以突破当前显存屏障。

你是一名 AI 行业编辑，请围绕下面这条热点输出一份资讯解读：
热点：Lightning Attention-2 新一代注意力机制要求：
1. 先用一句话解释这条热点在讲什么
2. 再总结它为什么重要
3. 说明会影响哪些 AI 产品或内容方向
4. 最后给出 3 个适合资讯站使用的标题