Scaling TransNormer to 175 Billion Parametes

线性注意力的Transformer大模型
2023

Transformer 存在局限。首要的一点，它们有着对于序列长度的二次时间复杂度，这会限制它们的可扩展性并拖累训练和推理阶段的计算资源和时间效率。 TransNormerLLM 是首个基于线性注意力的 LLM。
其中值得格外注意的一项改进是将 TransNormer 的 DiagAttention 替换成线性注意力，从而可提升全局的互动性能。研究者还引入了带指数衰减的 LRPE 来解决 dilution 问题。此外，研究者还**引入了 Lightning Attention（闪电注意力）**这种全新技术，并表示其可以将线性注意力在训练时的速度提升两倍，并且其还能通过感知 IO 将内存用量减少 4 倍。不仅如此，他们还简化了 GLU 和归一化方法，而后者将整体的速度提升了 20%。他们还提出了一种稳健的推理算法，可以在不同的序列长度下保证数值稳定和恒定的推理速度，由此能提升模型在训练和推理阶段的效率。

改进一：位置编码
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。

整体结构

在该结构中，输入 X 的更新通过两个连续步骤完成：首先，其通过使用了 SRMSNorm 归一化的门控式线性注意力（GLA）模块。然后，再次通过使用了 SRMSNorm 归一化的简单门控式线性单元（SGLU）模块。这种整体架构有助于提升模型的性能表现。下方给出了这个整体流程的伪代码：
在这里插入图片描述