科学家提出新一代注意力机制Lightning Attention，已在大语言模型上成功验证

来源：DeepTech深科技

‍

一直以来，减少 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）

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

图 | 钟怡然（来源：钟怡然）

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

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

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

参考资料：

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

运营/排版：何晨龙