放弃Softmax，首个线性注意力Transformer大模型TransNormerLLM：1750亿参数，速度、精度更优近日，上海人工智能实验室和OpenNLPLab的研究团队提出了一种名为TransNormerLLM的新型大型语言模型。该模型完全抛弃了传统的基于Softmax的注意力机制，转而采用了新提出的线性注意力机制。这一创新使得TransNormerLLM在准确度和效率方面均优于传统的Softmax注意力模型。一、模型背景与优势大型语言模型已广泛应用于自然语言处理（NLP）领域，显著提升了语言理解、生成和交互的能力。传统的Transformer架构，如GPT系列、BERT和BART等，均依赖于Softmax注意力机制来识别输入token之间的依赖关系，并感知全局位置。然而，这种机制存在局限性，尤其是序列长度的二次时间复杂度，限制了模型的可扩展性，并影响了训练和推理阶段的效率。为了克服这一局限，研究者们提出了多种序列建模方法，旨在将二次时间复杂度简化为线性。然而，这些方法在语言建模任务上的表现往往不尽如人意，且在真实场景中并未展现出明显的速度优势。TransNormerLLM的提出，则打破了这一僵局，成为首个基于线性注意力的大型语言模型，并在准确度和效率上实现了双重提升。二、模型架构与改进TransNormerLLM的构建基于之前的线性注意力架构TransNormer，并在此基础上进行了多项关键性改进，包括位置嵌入、线性注意力加速、门控机制、张量归一化和推理加速。位置编码：研究者为TransNormerLLM引入了带指数衰减的线性化相对位置编码（LRPE-d），以解决较低层中token之间缺乏全局互动能力的问题。门控机制：为了增强模型性能并使训练过程平滑，研究者采用了来自《Transformer quality in linear time》论文的Flash方法，并在token混合中使用了门控式线性注意力（GLA）结构。此外，还提出了Simple GLU（SGLU），去除了原始GLU结构的激活函数，因为门本身就能引入非线性。张量归一化：研究者使用了TransNormer中引入的NormAttention，并在TransNormerLLM中用新的简单归一化函数SimpleRMSNorm（SRMSNorm）替换了RMSNorm。三、整体结构与训练优化TransNormerLLM的整体结构如图1所示，输入X的更新通过两个连续步骤完成：首先通过使用了SRMSNorm归一化的门控式线性注意力（GLA）模块，然后再次通过使用了SRMSNorm归一化的简单门控式线性单元（SGLU）模块。这种整体架构有助于提升模型的性能表现。在训练优化方面，研究者引入了闪电注意力（Lightning Attention）算法，以加快注意力计算速度，并使新提出的线性注意力更适合IO处理。此外，还采用了全分片数据并行（FSDP）策略来优化内存利用率，以及激活检查点技术来减少内存中的激活数量，提高计算效率。同时，还使用了自动混合精度（AMP）来进一步节省资源并加快计算速度。四、语料库与实验为了验证TransNormerLLM的效果，研究者收集了一个大型语料库，总大小超过6TB，token数超过2万亿。在实验中，他们基于新的大型语料库进行了全面实验和控制变量研究，结果表明新方法的性能优于基于Softmax注意力的方法，并且具有更快的训练和推理速度。具体来说，在相同的配置下，当模型参数数量为385M和1B时，TransNormerLLM的性能比传统的Transformer分别好5%和9%。此外，通过一系列消融实验，研究者还验证了各项改进措施的有效性。五、系统优化与开源除了上述架构上的改进外，研究者还通过对线性transformer执行模型并行化而进行了系统工程优化。他们借鉴了英伟达的Megatron-LM模型并行化策略，将TransNormerLLM的两个主要模块SGLU和GLA的模型并行化分开执行。这一优化策略进一步提高了模型的训练效率。此外，为了促进LLM领域的研究和发展，上海人工智能实验室和OpenNLPLab的研究者将开源TransNormerLLM的预训练模型。他们表示，这一举措旨在让研究者和实践者能够基于他们的成果构建应用并探索用于LLM的高效Transformer结构。六、结论综上所述，TransNormerLLM作为一种新型的大型语言模型，在放弃Softmax注意力机制后，通过采用线性注意力机制并进行多项关键性改进，实现了在准确度和效率上的双重提升。这一创新不仅为NLP领域带来了新的突破和发展机遇，也为未来的大型语言模型研究提供了新的思路和方法。