【大模型系列篇】Transformer模型优化技巧和选择

木亦汐丫

已于 2024-08-30 09:42:07 修改

阅读量1k

点赞数 19

分类专栏：大模型文章标签： transformer RMSNorm SwiGLU GeGLU RoPE ALiBi GQA

于 2024-08-30 06:00:00 首次发布

本文链接：https://blog.csdn.net/Jackie_vip/article/details/141690113

版权

大模型专栏收录该内容

10 篇文章 0 订阅

订阅专栏

自从 Transformer 模型发布以来，研究人员针对训练稳定性、性能与计算效率提升等方面提出了多种改进方法。包括归一化、位置编码、激活函数和注意力机制。

`归一化`

`归一化方法`

大语言模型的预训练过程中经常会出现不稳定的问题。为了应对这一问题，深度学习方法通常会采用特定的归一化策略来加强神经网络训练过程的稳定性。

原始的 Transformer 模型主要使用了层归一化方法（Layer Normalization, LN）[1] 。

RMSNorm，为了提高层归一化的训练速度，RMSNorm [2] 仅利用激活值总和的均方根 RMS(x) 对激活值进行重新缩放。使用 RMSNorm的Transformer 模型相比于之前 LayerNorm 训练的模型在训练速度和性能上均具有一定优势。

DeepNorm，由微软的研究人员提出 [3]，旨在稳定深层 Transformer 的训练。具体而言，DeepNorm 在LayerNorm的基础上，在残差连接中对之前的激活值x按照一定比例 𝛼 进行放缩，可以有效提升了模型性能与训练稳定性。

`归一化模块位置`

除了归一化方法外，归一化模块的位置也具有重要的影响。归一化模块的位置通常采用前置或者后置。

层后归一化（Post-Layer Normalization, Post-Norm）

原始 Transformer 模型中所使用的一种归一化技术是Post-Norm。即归一化模块被放置于残差计算之后。计算公式如下：

理论上后向归一化具有很多优势。首先，有助于加快神经网络的训练收敛速度，使模型可以更有效地传播梯度，从而减少训练时间。其次，后向归一化可以降低神经网络对于超参数（如学习率、初始化参数等）的敏感性，使得网络更容易调优，并减少了超参数调整的难度。

然而，由于在输出层附近存在梯度较大的问题，采用 Post-Norm 的Transformer模型在训练过程中通常会出现不稳定的现象 [4]。

现有的大语言模型中，Post-Norm很少被单独使用，通常是与其他策略相结合应用。例如，GLM-130B 将 Post-Norm与DeepNorm结合使用。

层前归一化（Pre-Layer Normalization, Pre-Norm）

相较于 Post-Norm，Pre-Norm直接把每个子层加在了归一化模块之后，仅仅对输入的表示进行了归一化，从而可以防止模型的梯度爆炸或者梯度消失现象。当然，Pre-Norm在最后一个 Transformer 层后还额外添加了一个LayerNorm。计算公式如下：

虽然使用了Pre-Norm的Transformer模型在训练过程中更加稳定，但是性能却逊色于采用了Post-Norm的模型。尽管对于性能有一定的影响，但由于其能够有效维持训练的稳定性，很多主流的大语言模型仍然采用 Pre-Norm。

`激活函数`

前馈网络中激活函数的选择对于大语言模型的表现至关重要。通常来说，激活函数主要是为神经网络中引入非线性变化，从而提升神经网络的模型能力。在原始的Transformer 中采用了 ReLU激活函数。

该激活函数计算较为简单，仅仅是将对输入中每个神经元和“零值”进行比较，并将小于零的神经元的值设置为 0。然而，ReLU 可能会产生神经元失效的问题，被置为0的神经元将学习不到有用的信息。

ReLU 函数的具体形式如下所示：

针对ReLU存在的不足，研究人员进一步探索了ReLU函数的变种，以实现更好的性能。Swish 激活函数将神经元和该神经元的sigmoid激活的乘积作为新的激活函数。而GELU[5] 则利用标准高斯累积分布函数作为激活函数，被很多的Transformer模型所采用。

相比于原始的ReLU函数，这些新的激活函数通常能够带来更好的性能并且收敛性更好，但是计算过程更为复杂。Swish和GELU的数学表示如下：

近来，大语言模型（例如 PaLM 和 LaMDA）也经常采用 GLU激活函数以及它的变种 [6]，特别是SwiGLU和GeGLU。

不同于其他激活函数，GLU激活函数引入了两个不同的线性层。其中一个线性层的输出将被输入到一个激活函数（例如，GeGLU 采用 GELU 激活函数）中，其结果将和另一个线性层的输出进行逐元素相乘作为最终的输出。

相比于其他的激活函数，使用GLU激活函数变体通常能够带来更佳的性能表现 [7]。SwiGLU和 GeGLU激活函数的计算公式如下所示：

可以看到，目前大部分主流模型采用的是SwiGLU 或 GeGLU 激活函数。

`位置编码`

由于Transformer 模型中自注意力模块具有置换不变性，因此仅使用注意力机制无法捕捉序列中的顺序关系，从而退化为“词袋模型”。为了解决这一问题，需要引入位置编码（Position Embedding, PE）对于序列信息进行精确建模，从而将绝对或相对位置信息整合到模型中。

绝对位置编码，在原始的 Transformer 模型中，为了处理序列数据的顺序信息，采用了绝对位置编码方法。在编码器和解码器的输入端，根据输入的词元在序列中的绝对位置生成唯一的位置嵌入，并与词元的嵌入表示进行相加来注入位置信息。绝对位置编码公式如下:

其中， $p_t$ 表示位置 𝑡 的位置嵌入， $v_t$ 是该位置词元对应的词向量。

原始的Transformer采用了正余弦位置编码。该位置编码在不同维度上预先定义了特定的正弦或余弦函数，通过将词元的绝对位置作为输入代入这些函数，从而为这些维度赋予相应的值。对于维度大小为 𝐻 的位置嵌入，其第 𝑖 ∈ {1, . . . , 𝐻} 维的值按照如下方法进行设置：

此外，绝对位置编码还可以采用可学习的嵌入表示，并被很多早期的预训练语言模型（如 BERT）广泛采用。

相对位置编码，与绝对位置编码不同，相对位置编码是根据键和查询之间的偏移量计算得来的。计算得到的相对位置编码通常应用于注意力矩阵的计算中，而不是直接与词元本身的位置编码进行相加。

其中，Transformer-XL [8] 提出了一种相对位置编码方法，在计算键和查询之间的注意力分数时引入了相对位置信息。

对于使用绝对位置编码的模型，其注意力值可以进行进一步的分解：

而 Transformer-XL 对上述注意力值进行了改写，使用相对位置信息代替绝对位置信息。其公式表示如下所示（这里使用不同颜色与原始项进行对应）:

其中， $x_i$ 是每个词元对应的词向量（对应没有显式加入位置编码的词向量 $v_t$ ），而 $r_{i-j}$ 表示相对位置编码，𝒖 和 𝒗 是两个可学习的表示全局信息的参数。相比于绝对位置编码，注意力值的第二项中和第四项键对应的绝对位置编码 $W^{KT}p_j$ 被替换为相对位置编码 $r_{j}$ ，以引入相对位置信息；而第三和第四项中则使用全局参数 𝒖 和 𝒗 替换查询对应的绝对位置编码 $p_iW^Q$ ，用于衡量键的语义信息和相对位置信息本身的重要程度。

作为另一种方法，T5 提出了一种较为简化的相对位置编码。具体来说，它在注意力分数中引入了可学习的标量，这些标量是基于查询和键的位置之间的距离计算的。与绝对位置编码相比，应用了相对位置编码的 Transformer 模型常常可以对比训练序列更长的序列进行建模，即具备一定的外推能力 [9]。 T5 相对位置编码的计算可以表达为：

其中 $r_{i-j}$ 表示基于查询和键之间偏移的可学习标量。

旋转位置编码，RoPE 巧妙地使用了基于绝对位置信息的旋转矩阵来表示注意力中的相对位置信息。RoPE 根据位置信息为序列中每个词元所对应的设置了独有的旋转矩阵，并和对应的查询和键进行相乘进行融合。位置索引为 𝑡 对应的旋转矩阵定义如下所示：

利用旋转矩阵中三角函数的特性，位置索引为 𝑖 的旋转矩阵和位置索引为 𝑗 的旋转矩阵的转置的乘积等同于位置索引为它们相对距离𝑖-𝑗 的旋转矩阵，即 $R_{\theta ,i}R^T_{\theta , j} = R_{\theta ,i-j}$ 。通过这种方式，键和查询之间的注意力分数能够有效融入相对位置信息。注意力矩阵的公式可以进一步变为如下形式：

根据旋转矩阵的定义，RoPE 在处理查询和键向量的时候，将每对连续出现的两个元素视为一个子空间。因此，对于一个长度为 𝐻 的向量来说，将会形成 𝐻/2 个这样的子空间。在这些子空间中，每一个子空间 𝑖 ∈ {1, . . . , 𝐻/2} 所对应的两个元素都会根据一个特定的旋转角度 𝑡 · 𝜃𝑖 进行旋转，其中 𝑡 代表位置索引，而 𝜃𝑖 表示该子空间中的基。与正弦位置嵌入类似，RoPE 将基 𝜃𝑖 定义为底数 𝑏（默认值是 10000）的指数：

进一步，每个子空间定义了波长 𝜆𝑖，即在该子空间上完成一个完整周期（2𝜋）旋转所需的距离：

由于 RoPE 具有良好的性能以及长期衰减的特性，已经主流的大语言模型广泛采用。

ALiBi位置编码，ALiBi [9] 是一种特殊的相对位置编码，主要用于增强Transformer模型的外推能力。具体来说，ALiBi 通过在键和查询之间的距离上施加相对距离相关的惩罚来调整注意力分数。其计算公式如下：

其中，𝑖−𝑗 是查询和键之间的位置偏移量，𝑚 是每个注意力头独有的惩罚系数。与T5等模型中的相对位置编码不同，ALiBi 中的惩罚分数是预先设定的，不需要引入任何可训练的参数。此外，ALiBi 展现出了优秀的外推性能，能够对于超过上下文窗口更远距离的词元进行有效建模。

`注意力机制`

注意力机制是 Transformer 架构中的核心技术，它能够针对序列中的词元对构建交互关系，聚合来自于不同位置的语义信息。下面介绍四种常见的注意力机制的设计方法。

完整自注意力机制，原始的 Transformer 模型中，注意力机制通过成对的方式进行序列数据的语义建模，充分考虑了序列中所有词元之间的相互关系。其中，每个词元在注意力计算中都需要对于其前序的所有词元的键值对予以访问，因此对于序列长度为 𝑇 的序列需要 𝑂(𝑇 2 ) 的计算复杂度。此外，Transformer 还引入了多头注意力机制，将查询、键和值在不同的语义空间进行线性投影，然后将每个头的输出进行聚合形成最终的输出。完整自注意力可以阅读原论文解读篇《论文解读:Transformer - Attention Is All You Need》。

稀疏注意力机制，尽管完整自注意力机制具有较强的建模能力，但是它需要平方级的计算复杂性。在处理长序列时较为显著，带来了较大的计算和存储开销。为了降低注意力机制的计算复杂度，研究人员提出了多种高效的注意力变种。其中，滑动窗口注意力机制（Sliding Window Attention, SWA）是大语言模型中使用最多的一种稀疏注意力机制。不同于完整的注意力机制，滑动窗口注意力根据词元位置，仅仅将位置索引上距离该词元一定范围内的词元考虑到注意力的计算中。具体来说，滑动窗口注意力设置了一个大小为 𝑤 的窗口，对每个词元 𝑢𝑡，只对窗口内的词元 [𝑢𝑡−𝑤+1, . . . , 𝑢𝑡] 进行注意力计算，从而将复杂度降低到 𝑂(𝑤𝑇)。进一步，通过信息的逐层传递，模型实现了随着层数线性增长的感受野，从而获取远处词元的信息。更多细节可以参考复旦大学邱锡鹏教授的Transformers综述中关于稀疏注意力机制的介绍《Transformers综述--邱锡鹏》。

多查询/分组查询注意力，为了提升注意力机制的效率，多查询注意力（MultiQuery Attention, MQA）提出针对不同的头共享相同的键和值变换矩阵 [10]。这种方法减少了访存量，提高了计算强度，从而实现了更快的推理速度。并且对于模型性能产生的影响也比较小。

为了结合多查询注意力机制的效率与多头注意力机制的性能，研究人员进一步提出了分组查询注意力机制（Grouped-Query Attention, GQA）[11]。GQA 将全部的头划分为若干组，并且针对同一组内的头共享相同的变换矩阵。这种注意力机制有效地平衡了效率和性能，被 LLaMA-2 模型所使用。

硬件优化的注意力机制，除了在算法层面上提升注意力机制的计算效率，还可以进一步利用硬件设施来优化注意力模块的速度和内存消耗。其中，两个具有代表性的工作是 FlashAttention [12] 与 PagedAttention [13]。相比于传统的注意力实现方式，FlashAttention 通过矩阵分块计算以及减少内存读写次数的方式，提高注意力分数的计算效率；PagedAttention 则针对增量解码阶段，对于 KV 缓存进行分块存储，并优化了计算方式，增大了并行计算度，从而提高了计算效率。关于FlashAttention和PagedAttention 可阅读《自注意力机制计算加速工程优化技巧》。

`总结`

综合以上我们分析的归一化、位置编码、激活函数和注意力机制，在模型配置上给出一些推荐建议。为了增强模型的训练稳定性，建议采用前置的RMSNorm作为层归一化方法。同时在选择激活函数时，为了获得更优的模型性能，可以优先考虑使用 SwiGLU 或 GeGLU。而对于位置编码，可以优先选择 RoPE 或者 ALiBi，这两种位置编码方法在建模长序列数据时通常能够具有较好的性能。至于注意力机制的选择需要综合模型的效率和性能两方面来考量。

`参考文献`

[1] Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. “Layer Normalization”. In: arXiv preprint arXiv:abs/1607.06450 (2016).

[2] Biao Zhang and Rico Sennrich. “Root Mean Square Layer Normalization”. In: NeurIPS. 2019.

[3] Hongyu Wang et al. “DeepNet: Scaling Transformers to 1, 000 Layers”. In: arXiv preprint arXiv:abs/2203.00555 (2022).

[4] Aohan Zeng et al. “GLM-130B: An Open Bilingual Pre-trained Model”. In: arXiv preprint arXiv:abs/2210.02414 (2022).

[5] Ruibin Xiong et al. “On layer normalization in the transformer architecture”. In: ICML. 2020.

[6] Dan Hendrycks and Kevin Gimpel. “Gaussian error linear units (gelus)”. In: arXiv preprint arXiv:1606.08415 (2016).

[7] Noam Shazeer. “GLU Variants Improve Transformer”. In: arXiv preprint arXiv:2002.05 202 (2020)

[8] Sharan Narang et al. “Do Transformer Modifications Transfer Across Implementations and Applications?” In: EMNLP. 2021.

[9] Zihang Dai et al. “Transformer-XL: Attentive Language Models beyond a Fixed-Length Context”. In: ACL. 2019.

[10] Ofir Press, Noah A. Smith, and Mike Lewis. “Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation”. In: ICLR. 2022.

[11] Noam Shazeer. “Fast Transformer Decoding: One Write-Head is All You Need”. In: arXiv preprint arXiv:1911.02150 (2019).

[12] Joshua Ainslie et al. “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”. In: arXiv preprint arXiv:2305.13245 (2023)

[13] Tri Dao et al. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness”. In: NeurIPS. 2022.

[14] Woosuk Kwon et al. “Efficient memory management for large language model serving with pagedattention”. In: SOSP. 2023.

木亦汐丫

关注

19
点赞
踩
21

收藏

觉得还不错? 一键收藏
0
评论
【大模型系列篇】Transformer模型优化技巧和选择

自从 Transformer 模型发布以来，研究人员针对训练稳定性、性能与计算效率提升等方面提出了多种改进方法。包括归一化、位置编码、激活函数和注意力机制。为了增强模型的训练稳定性，建议采用前置的RMSNorm作为层归一化方法。同时在选择激活函数时，为了获得更优的模型性能，可以优先考虑使用 SwiGLU 或 GeGLU。而对于位置编码，可以优先选择 RoPE 或者 ALiBi，这两种位置编码方法在建模长序列数据时通常能够具有较好的性能。至于注意力机制的选择需要综合模型的效率和性能两方面来考量。
复制链接

扫一扫

专栏目录