GPT 等大型语言模型（LLM）的成功离不开 Softmax 注意力机制，但这一机制也存在着成本高等一些缺点。
近日，上海人工智能实验室和 OpenNLPLab 的一个研究团队提出了一种新的大型语言模型 TransNormerLLM，其中完全抛弃了基于 Softmax 的注意力机制，而是使用了新提出的线性注意力。据介绍，TransNormerLLM 是首个基于线性注意力的大型语言模型（LLM），其在准确度和效率方面的表现优于传统的基于 Softmax 注意力的模型。研究者也将发布其预训练模型的开源版本。
大型语言模型已经为自然语言处理（NLP）领域带来了变革。它们能出色地处理许多不同类型的任务，提升了在计算框架中使用人类语言进行理解、生成和交互的能力。之前语言建模的发展主要围绕 Transformer 架构，其中堪称支柱的模型包括基本的 Transformer、GPT 系列、BERT 和 BART 等开创性的模型。Transformer 架构的成功立足于 softmax 注意力机制，该机制可以识别出数据驱动的模型方案中输入 token 之间的依赖关系，其还能感知全局位置，让模型能有效处理自然语言的长程动态情况。
即便如此，传统的 Transformer 依然存在局限。首要的一点，它们有着对于序列长度的二次时间复杂度，这会限制它们的可扩展性并拖累训练和推理阶段的计算资源和时间效率。为了将这种二次时间复杂性简化至线性复杂性，已经有不少研究者提出了多种不同的序列建模方法。但是，这些方法却难以用于 LLM，原因有二：1) 它们在语言建模任务上的表现往往不如人意；2) 它们在真实世界场景中并没有展现出速度优势。
这篇论文提出的 TransNormerLLM 是首个基于线性注意力的 LLM，其在准确度和效率方面都优于传统的 softmax 注意力。TransNormerLLM 的构建基于之前的线性注意力架构 TransNormer，同时也做了一些修改让性能更优。TransNormerLLM 中关键性的改进包括位置嵌入、线性注意力加速、门控机制、张量归一化和推理加速。
其中值得格外注意的一项改进是将 TransNormer 的 DiagAttention 替换成线性注意力，从而可提升全局的互动性能。研究者还引入了带指数衰减的 LRPE 来解决 dilution 问题。此外，研究者还引入了 Lightning Attention（闪电注意力）这种全新技术，并表示其可以将线性注意力在训练时的速度提升两倍，并且其还能通过感知 IO 将内存用量减少 4 倍。不仅如此，他们还简化了 GLU 和归一化方法，而后者将整体的速度提升了 20%。他们还提出了一种稳健的推理算法，可以在不同的序列长度下保证数值稳定和恒定的推理速度，由此能提升模型在训练和推理阶段的效率。
为了验证 TransNormerLLM 的效果，研究者精心收集了一个大型语料库，其大小超过 6TB，token 数更是超过 2 万亿。为了确保数据的质量，他们还开发了一种用于过滤所收集语料库的自清理（self-cleaning）策略。如表 1 所示，研究者对原始 TransNormer 模型进行了扩展，得到了参数量从 3.85 亿到 1750 亿的多个 TransNormerLLM 模型。然后他们基于新的大型语料库进行了全面实验和控制变量研究，结果表明新方法的性能优于基于 softmax 注意力的方法并且还有更快的训练和推理速度。
表 1：TransNormerLLM 模型的多种变体
为了促进 LLM 领域的研究，上海人工智能实验室和 OpenNLPLab 的这些研究者也将开源自己的预训练模型，他们表示这是为了「让研究者和实践者基于我们的成果构建应用以及探索用于 LLM 的高效 transformer 结构」。
TransNormerLLM
架构改进
下面将简单介绍 TransNormerLLM 的各个模块以及研究者提出的一些改进措施。
改进一：位置编码
TransNormer 中的较低层使用了 DiagAttention 来避免 dilution 问题。但是，这会导致 token 之间缺乏全局互动能力。为了解决这个问题，研究者为 TransNormerLLM 使用了带指数衰减的 LRPE（线性化相对位置编码），从而可在较低层保留完整的注意力。研究者把这种方法称为 LRPE-d。
改进二：门控机制
门控可以增强模型的性能并使训练过程平滑。研究者为 TransNormerLLM 使用了来自论文《Transformer quality in linear time》的 Flash 方法并在 token 混合中使用了门控式线性注意力（GLA）的结构。
为了进一步提升模型速度，他们还提出了 Simple GLU（SGLU），其去除了原始 GLU 结构的激活函数，因为门本身就能引入非线性。
改进三：张量归一化
研究者使用了 TransNormer 中引入的 NormAttention。在 TransNormerLLM 中，他们使用一种新的简单归一化函数 SimpleRMSNorm（简写为 SRMSNorm）替换了 RMSNorm。
整体结构
图 1 展示了 TransNormerLLM 的整体结构。
图 1：新提出模型的整体架构
在该结构中，输入 X 的更新通过两个连续步骤完成：首先，其通过使用了 SRMSNorm 归一化的门控式线性注意力（GLA）模块。然后，再次通过使用了 SRMSNorm 归一化的简单门控式线性单元（SGLU）模块。这种整体架构有助于提升模型的性能表现。下方给出了这个整体流程的伪代码：
训练优化
闪电注意力
为了加快注意力计算速度，研究者引入了闪电注意力（Lightning Attention）算法，这能让新提出的线性注意力更适合 IO（输入和输出）处理。
算法 1 展示了闪电注意力的前向通过的实现细节，算法 2 则是后向通过的。研究者表示，他们还有一个可以更快计算梯度的实现版本，这会在未来发布。
模型并行化
为了在计算机集群上分散所有模型参数、梯度和优化器状态张量，研究者使用了全分片数据并行（FSDP/Fully Sharded Data Parallel）。这种策略性分区方法可减少对每个 GPU 的内存占用，从而优化了内存利用率。为了进一步提高效率，他们使用了激活检查点（Activation Checkpointing），这可减少后向通过过程中缓存在内存中的激活数量。相反，当计算这些梯度时，这些梯度会被移除并重新计算。该技术有助于提升计算效率和节省资源。此外，为了在减少 GPU 内存消耗的同时加快计算速度，研究者还使用了自动混合精度（AMP）。
除了上述成果外，研究者还更进一步通过对线性 transformer 执行模型并行化而进行了系统工程优化，其灵感很大程度上来自于英伟达的 Megatron-LM 模型并行化，在传统的 Transformer 模型中，每个 transformer 层都有一个自注意力模块，其后跟着一个两层多层感知器（MLP）模块。当使用 Megatron-LM 模型并行性时，是在这两个模块上独立使用。类似地，TransNormerLLM 结构也是由两个主要模块构成：SGLU 和 GLA；这两者的模型并行化分开执行。
稳健推理
这让 TransNormerLLM 能以 RNN 的形式执行推理。算法 3 给出了这个过程的细节。但其中存在数值精度问题。
为了避免这些问题，研究者提出了稳健推理算法，见算法 4。
原推理算法和稳健推理算法得到的结果是一样的。
语料库
研究者从互联网收集了大量可公开使用的文本，总大小超过 700TB。收集到的数据经由他们的数据预处理程序进行处理，如图 2 所示，留下 6TB 的干净语料库，其中包含大约 2 万亿 token。为了提供更好的透明度，帮助用户更好理解，他们对数据源进行了分门别类。表 2 给出了具体的类别情况。
图 2：数据预处理流程
表 2：语料库统计数据
实验
研究者在 Metaseq 框架中使用 PyTorch 和 Trition 实现了 TransNormerLLM。模型的训练使用了 Adam 优化器，FSDP 也被用于高效地将模型扩展到 NVIDIA A100 80G 集群。他们也适当地使用了模型并行技术来优化性能。
架构消融实验
表 3：Transformer vs TransNormerLLM。在相同的配置下，当模型参数数量为 385M 和 1B 时，TransNormerLLM 的性能比 Transformer 分别好 5% 和 9%。
表 4：TransNormer vs TransNormerLLM。TransNormerLLM 的实验结果是最好的。
表 5：位置编码组合使用 LRPE+LRPE-d 得到的结果最优。
表 6：衰减温度方面的消融实验结果。结果表明新方法更优。
表 7：门控机制方面的消融实验结果。使用该门控机制的模型表现更好。
表 8：GLA 激活函数的消融实验结果。用不同激活函数得到结果差不多。
表 9：GLU 激活函数的消融实验结果。去除激活函数对结果没有负面影响。
表 10：归一化函数。使用以下归一化函数的结果差异不大。
图 3：对 SRMSNorm 实现的性能评估
图 4：线性注意力和闪电注意力之间的内存和速度比较
图 5：推理时间和内存占用情况
系统优化
表 11：模型并行性性能
表 12：训练不同大小的模型的效率
表 13：训练 Transformer 和 TransNormerLLM 的最大上下文长度