首个线性注意力Transformer大模型：1750亿参数，速度、精度更优 _大模型

编辑：Panda

GPT 等大型语言模型（LLM）的成功离不开 Softmax 注意力机制，但这一机制也存在着成本高等一些缺点。

近日，上海人工智能实验室和 OpenNLPLab 的一个研究团队提出了一种新的大型语言模型 TransNormerLLM ，其中完全抛弃了基于 Softmax 的注意力机制，而是使用了新提出的线性注意力。据介绍， TransNormerLLM 是首个基于线性注意力的大型语言模型（LLM），其在准确度和效率方面的表现优于传统的基于 Softmax 注意力的模型。研究者也将发布其预训练模型的开源版本。

文章插图
论文：
https://arxiv.org/abs/2307.14995
模型：
https://Github.com/OpenNLPLab/TransnormerLLM
大型语言模型已经为自然语言处理（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 。