一直以来，减少 Transformer 的二次计算复杂度都是一个老生常谈的问题。

当前算力的高速增长（V100-A100-H100-GH200）基本覆盖了其二次计算复杂度带来的算力需求，使得目前工业界对于解决 Transformer 二次计算复杂度的需求并不强烈。

同时，当前的线性解决方案仍停留在研究阶段，最终效果和实际效率并没有得到广泛的验证，导致工业界仍以 Transformer 架构作为首选技术方案。

然而，算力的增长终归会遇到瓶颈，设计出比 Transformer 更高效的架构来取代 Transformer 基本会成为历史的必然。

近日，一支团队提出了新一代注意力机制 Lightning Attention，针对性地设计了新的网络基础架构 TransNormerLLM（TNL）。

其在计算效率和计算效果上均能大幅超越最新的 Transformer 架构，并在大语言模型上验证了它的可行性。

Lightning Attention 与传统的 Transformer 注意力机制有着本质上的不同，它是一个基于线性注意力的机制。

通过交换 QKV 相乘的先后顺序（即 Q（KV）），线性注意力的理论计算复杂度与序列长度呈线性增长的关系。

需要注意的是，线性注意力与 Transformer 虽然在同年推出，且有着理论复杂度低的优势，但却始终无法成为主流方案。

导致这种现象的原因有两个：

其一，效果差。

相比于传统的基于 softmax 的注意力（softmax attention），线性注意力方案在标准学术数据集上有着明显的效果差异。

其二，训练速度慢。

线性注意力为了实现线性计算复杂度，在因果语言建模中需要一个名叫累积求和（cumsum）的操作。

这个操作是通过循环实现的，对于 GPU（graphics processing unit，图形处理器）这种并行架构不友好，效率不高，在实际实现中往往远慢于直接左乘的矩阵乘法（（QK）V）。

因此，由于线性注意力速度也不快，效果也不好，就成为了中看不中用的“银样蜡枪头”。

然而，在本次研究人员推出的 Lightning Attention 和 TNL 中，这两个问题均被解决。

在训练速度方面，他们提出了 Lightning Attention，这是一种新的线性注意力的实现，首次让线性注意力在实际应用中实现其理论线性计算优势。

在 Lightning Attention 中，他们采用分而治之的策略来计算注意力，将计算分为块内和块间。

其中，块内计算采用左乘的形式，块间则采用右乘的形式来绕开 cumsum 操作。

同时，他们针对 GPU 进行了类似于 FlashAttention 的 IO 优化，最终 Lightning Attention 实现了在固定显存的条件下，训练速度随着序列长度的增加保持恒定的效果。

（来源：arXiv）

在效果方面，他们针对 Lightning Attention 推出了一个新的网络框架 TNL。

TNL 基于该团队于 2022 年 EMNLP 推出的 TransNormer 架构打造而来，通过调整位置编码、归一化方式、以及添加门控系统，TNL 的速度更快、效果更好、训练更稳定。

从下图可以看出，TNL 在 1B 和 3B 模型大小上的效果远好于同期的 RNN（HGRN）和长卷积（TNN）建模方案，同时也优于最先进的基于 Transformer 的架构 LLaMA。

（来源：arXiv）

研究人员表示：“这项研究成果是一个革命性的技术，打破了序列长度对于大语言模型甚至多模态大模型的限制，使得处理无限序列长度成为可能。”

相比于现有的 Transformer 架构在序列长度增加时，所需要的计算资源呈几何倍数增长，TNL 无需增加任何开销即可任意扩增序列长度。

在未来的超长多轮对话、高分辨率图像生成、长视频生成等领域，TNL 将在不损失精度的情况下，极大地减少计算量，并能在有限的计算资源下，促进 AIGC 技术的发展。

当然，提出一个新的注意力机制或者一套网络框架是简单的，困难的是如何说服别人采用或者支持这个新架构。

线性注意力机制因为提出时间早，在很长的一段时间里都无法在实际应用中证明自己的高效性，人们对于这个新实现很容易抱着一种“狼来了”的心态，担心又是一个“漂亮的烟花”。

特别是当前有一个观点认为：这些高效序列建模方案比如 RWKV、HGRN、TNN 等，只能在网络参数较小的情况下效果跟 Transformer 差不多。 但是，在网络参数增大的情况下，效果会大幅降低。

为了证明 TNL 在大模型下的有效性，本次研究人员需要自己先训练出大模型出来，这需要投入大量的计算资源。

课题组表示：“在这里，非常感谢领导和兄弟部门的支持，让我们借调了大量的 GPU 资源。”

凭借这些帮助，他们才能先后在自有语料（1.4TB tokens）上训练了 7B 和 15B 大小的大语言模型，在标准 LLM Benchmark 上显示出了有竞争力的结果。

同时，为了让整个训练过程保持透明，他们首次通过直播训练的方式，让每一个人都可以随时且时时跟踪训练结果，并且获取中间的 checkpoints。

据他们所知，极少有公司会同步直播整个大模型训练过程，因为训练结果有很强的不确定性。

而该团队这样做的原因在于需要向大家证明这个网络的有效性，并且他们自身对于 TNL 也有着强烈的信心，相信不会在直播中翻车。

日前，相关论文以《Lightning Attention-2：在大型语言模型中处理无限序列长度的免费午餐》（）为题发在 arXiv[1]。

图 | 相关论文（来源：arXiv）

秦臻是第一作者，上海人工实验室青年科学家担任通讯作者。

图 | 钟怡然（来源：）

当前的 TNL 已经集成了模型并行，并在 175B 模型大小下进行了速度和显存测试，但是支持的序列长度仍局限于单个 GPU 的显存，无法发挥大规模 GPU 集群的优势。

在未来，他们将利用 Lightning Attention 的特性，推出针对线性注意力机制的序列并行方案，让无限序列长度真正成为可能。

另外，他们也将研究 Lightning Attention 的 encoder 架构，让它真正做到在各个领域都可以取代传统的 softmax attention。

参考资料：

1.https://arxiv.org/abs/2401.04658

运营/排版：何晨龙

01/

02/

03/

04/

05/