Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

7.2 Transformer 块

自注意力(self-attention)计算是所谓 变换器块(transformer block)的核心。除自注意力层之外,Transformer 块还包括另外三类层:(1)前馈层,(2)残差连接,以及(3)归一化层(通常俗称“层归一化”,即 layer norm)。

图 7.7 展示了一个 Transformer 块,并勾勒出一种称为 残差流(residual stream)的常见理解方式(Elhage et al., 2021)。在残差流视角下,我们把单个词元 i 通过 Transformer 块的处理过程,看作词元位置 i 上一条由 d 维表示构成的单一信息流。这条残差流从原始输入向量开始,各个组件从残差流中读取输入,再把自己的输出加回流中。

图 7.7 Transformer 块的架构及其残差流。大部分信息沿残差流向上传递,只有注意力模块会感知来自先前词元位置上其他残差流的信息。本图以及本章其余部分采用该架构的预归一化版本,即层归一化发生在注意力层和前馈层之前,而不是之后。

信息流底部的输入是一个维度为 dd 的词元嵌入。这个初始嵌入通过残差连接向上传递,并由 Transformer 的其他组件逐步将信息叠加到其中:一个是我们已经见过的注意力层,另一个是即将介绍的前馈层。在注意力层和前馈层之前,还会执行一种称为层归一化的计算。

因此,初始向量先经过层归一化和注意力层,所得结果再加回残差流;在这里,它被加到原始输入向量 xi\mathbf { x } _ { i } 上。随后,这个求和向量再次经过另一个层归一化和一个前馈层,它们的输出又被加回残差流。我们用 hi\mathbf { h } _ { i } 表示词元 i 经过 Transformer 块后得到的输出。

我们已经介绍了注意力层,下面将在处理词元位置 i 上单个输入 xi\mathbf { x } _ { i } 的语境中,介绍前馈计算和层归一化计算。

7.2.1 前馈层

前馈层(feedforward layer)是一个全连接的两层网络,也就是说,它有一个隐藏层和两个权重矩阵(第 6 章)。

前馈层是 逐位置的(position-wise),这意味着它独立作用于每个词元位置 i。这与注意力网络形成对照:注意力网络的任务是在不同词元位置之间混合信息。前馈层的权重在各个位置之间共享,也就是说,每个词元位置都应用相同的参数;但不同层所用的参数彼此不同。

通常会将前馈网络隐藏层的维度 dffd _ { \mathrm { f f } } 设为模型维度 dd 的 4 倍。(例如,在最初的 Transformer 模型中,d=512d = 512,而 dff=2048d _ { \mathrm { f f } } = 2048。)这意味着前馈网络占据了 Transformer 中的大多数参数,而且这些前馈网络参数似乎编码了 Transformer 中大部分事实知识(Geva et al., 2021; Meng et al., 2022)。

如今,大多数 LLM 实际上使用 门控前馈层(gated feedforward layer),其中采用一种称为 SwiGLU 的激活函数(Shazeer, 2020);它的效果优于 ReLU 等较简单的函数。SwiGLU 属于 门控线性单元(Gated Linear Unit, GLU)函数家族(Dauphin et al., 2017),这类函数对输入的两个线性变换执行逐分量乘积(\otimes)。其基本思想是:一个函数给出值,另一个函数充当门,指定允许其中多少值通过。SwiGLU 使用 Swish 函数(Ramachandran et al., 2017)作为门:

Swishβ(x)=xσ(βx)(7.21)\operatorname{Swish} _ {\beta} (\mathbf {x}) = \mathbf {x} \sigma (\beta \mathbf {x})\tag{7.21}

于是,SwiGLU 前馈计算的方程为:

FFNSwiGLU(x,W1,W2,W3,b1,b2,b3,β)=(Swishβ(xW1+b1)(xW2+b2))W3+b3(7.22)\begin{aligned} \operatorname{FFN} _ {\text {SwiGLU}} \left(\mathbf {x}, \mathbf {W} _ {1}, \mathbf {W} _ {2}, \mathbf {W} _ {3}, b _ {1}, b _ {2}, b _ {3}, \beta\right) \\ = \left(\operatorname{Swish} _ {\beta} \left(\mathbf {x} \mathbf {W} _ {1} + b _ {1}\right) \otimes \left(\mathbf {x} \mathbf {W} _ {2} + b _ {2}\right)\right) \mathbf {W} _ {3} + b _ {3}\tag{7.22} \end{aligned}

由于门控激活函数含有用于门控的额外参数,门控模型中的维度 dffd _ { \mathrm { f f } } 通常会设得略小于常规的 4d4d

7.2.2 层归一化

在 Transformer 块的两个阶段中,我们会对向量进行归一化(Ba et al., 2016)。这一过程称为 层归一化(layer normalization,简称 layer norm)。它是多种归一化方法之一;这些方法通过把隐藏层的取值维持在便于基于梯度进行训练的范围内,来改善深度神经网络的训练性能。

层归一化可以看作统计学中 z-score 的一种变体,它作用于隐藏层里的单个向量。也就是说,“层归一化”这个名称有些容易引起误解:它并不是应用于整个 Transformer 层,而只是应用于单个词元的嵌入向量。因此,层归一化的输入是一个维度为 dd 的向量,输出则是归一化后的同一个向量,维度仍为 dd。层归一化的第一步,是在待归一化向量的各个元素上计算均值 μ\mu 和标准差σ\sigma。给定一个维度为 dd 的嵌入向量 xx,这两个量按如下方式计算。

μ=1dj=1dx[j](7.23)\mu = \frac {1}{d} \sum_ {j = 1} ^ {d} x [ j ]\tag{7.23}
σ=1dj=1d(x[j]μ)2(7.24)\sigma = \sqrt {\frac {1}{d} \sum_ {j = 1} ^ {d} (x [ j ] - \mu) ^ {2}}\tag{7.24}

得到这些值后,从向量的每个分量中减去均值,再除以标准差,即可完成归一化。计算得到的新向量均值为零、标准差为一。

x^=(xμ)σ(7.25)\hat {\mathbf {x}} = \frac {(\mathbf {x} - \boldsymbol {\mu})}{\sigma}\tag{7.25}

最后,在层归一化的标准实现中,还会引入两个可学习参数 γ\gammaβ\beta,分别表示增益值和偏移值。

LayerNorm(x)=γ(xμ)σ+β(7.26)\operatorname{LayerNorm} (\mathbf {x}) = \gamma \frac {(\mathbf {x} - \mu)}{\sigma} + \beta\tag{7.26}

在实践中,许多现代网络使用一种更简单的归一化方法,称为 RMSNorm(Zhang and Sennrich, 2019)。它将向量除以均方根统计量来重新缩放向量,但省略了均值中心化步骤。

7.2.3 组合所有组件

我们可以用一个方程表示每项组件计算,从而分解 Transformer 块所计算的函数。下面使用 t(形状为 [1×d][ 1 \times d ])代表 Transformer,并用上标区分块内的每一步计算:

ti1=LayerNorm(xi)(7.27)\mathbf {t} _ {i} ^ {\mathbf {1}} = \operatorname{LayerNorm} (\mathbf {x} _ {i})\tag{7.27}
ti2= MultiHeadAttention (ti1,[t11,,ti11])(7.28)\mathbf {t} _ {i} ^ {2} = \text { MultiHeadAttention } (\mathbf {t} _ {i} ^ {1}, [ \mathbf {t} _ {1} ^ {1}, \dots , \mathbf {t} _ {i - 1} ^ {1} ])\tag{7.28}
ti3=ti2+xi(7.29)\mathbf {t} _ {i} ^ {3} = \mathbf {t} _ {i} ^ {2} + \mathbf {x} _ {i}\tag{7.29}
ti4= LayerNorm (ti3)(7.30)\mathbf {t} _ {i} ^ {4} = \text { LayerNorm } (\mathbf {t} _ {i} ^ {3})\tag{7.30}
ti5=FFN(ti4)(7.31)\mathbf {t} _ {i} ^ {\mathbf {5}} = \mathrm{FFN} (\mathbf {t} _ {i} ^ {\mathbf {4}})\tag{7.31}
hi=ti5+ti3(7.32)\mathbf {h} _ {i} = \mathbf {t} _ {i} ^ {5} + \mathbf {t} _ {i} ^ {3}\tag{7.32}

请注意,唯一会接收来自其他词元(其他残差流)信息的组件是多头注意力。正如式 7.28 所示,它会查看上下文中的所有相邻词元。不过,注意力的输出随后会被加到当前词元的嵌入流中。事实上,Elhage et al. (2021) 表明,我们可以把注意力头理解为确实在把信息从相邻词元的残差流移入当前残差流。因此,每个位置上的高维嵌入空间既包含当前词元的信息,也包含相邻词元的信息,尽管这些信息位于向量空间的不同子空间中。图 7.8 将这种移动过程可视化。由此,我们把注意力函数称为该架构中的 词元混合(token-mixing)组件,因为它把相邻词元流中的信息混合到当前信息流中。

至关重要的是,Transformer 块的输入和输出维度保持一致,因此可以堆叠。输入到块中的每个词元向量 xi\mathbf { x } _ { i } 的维度为 dd,输出 hi\mathbf { h } _ { i } 的维度同样为 dd。用于大语言模型的 Transformer 会堆叠许多这样的块:从 12 层(用于 T5 或 GPT-3-small 语言模型),到 96 层(用于 GPT-3 175B),较新的模型甚至会使用更多层。稍后我们还会回到这个堆叠问题。

图 7.8 一个注意力头可以把信息从词元 A 的残差流移入词元 B 的残差流。

式 7.27 及其后各式只描述了单个 Transformer 块,但残差流这一比喻贯穿所有 Transformer 层:以一个 12 层 Transformer 为例,它从第一个 Transformer 块一直延续到第 12 个块。在较低的 Transformer 块中,残差流表示当前词元;在最高的 Transformer 块中,残差流通常表示下一个词元,因为在整个网络最末端,模型接受的训练目标是预测下一词元。

堆叠许多个块以后,还需要满足一项要求:在最后一个(最高的)Transformer 块末端,需要对每条词元流最后的 hi\mathbf{h}_i 额外执行一次层归一化(位置就在我们稍后会定义的语言模型头层下方)。[2]

Footnotes
  1. 这是最常见的 Transformer 架构,即预归一化(prenorm)架构。Vaswani et al. (2017) 对 Transformer 的原始定义采用了另一种后归一化(postnorm)架构,其中层归一化位于注意力层和 FFN 层之后;将层归一化放在这些层之前效果更好,但确实需要在末端增加这一额外的层归一化。