第 8 章 · Transformer

“记忆的真正艺术，就是注意力的艺术。”
     Samuel Johnson, Idler #74, 1759 年 9 月

8   Transformer

本章介绍 transformer，也就是构建大型语言模型的标准架构。正如上一章所讨论的，基于 transformer 的大型语言模型已经彻底改变了语音与语言处理领域。事实上，本教材后续每一章都会用到它们。和上一章一样，本章重点讨论 transformer 在从左到右的语言建模中的用法；这种建模方式有时也称为因果语言建模或自回归语言建模。在这种任务中，模型得到一个输入词元序列，并在先前上下文的条件下逐个预测输出词元。

图 8.1 沿着单个词元在网络各层中向上传递的路径，概略展示了 transformer 架构。每个词元首先由嵌入矩阵 $E$ 转换为一个嵌入。回忆第 6 章第 6.5 节，$E$ 是一个线性层，它把词元 id 映射为表示该词元的向量嵌入。词汇表中的每个词元在 $E$ 中都有一个初始嵌入表示。

Transformer 还具有一种特殊机制，用来编码词元在输入字符串中的位置或索引；这个位置信息会直接加到词元嵌入上。这样得到的嵌入同时表示词和它的位置，然后被送入一组 $N$ 个 transformer 块。

一种常见的理解方式，是把每个 transformer 块看成某条流的一部分：输入嵌入会被直接向上传到输出，同时又被各种处理模块逐步丰富。这些模块包括多头注意力层、前馈网络和层归一化。某一层上这条流的值，是原始嵌入与此前所有层、所有块输出之和。

Transformer 的核心直觉，也是它区别于第 6 章中前馈层的组件，是多头注意力层，也称为自注意力层。注意力可以看作一种机制：通过关注并整合周围词元的信息，构建某个词元含义的上下文表示，从而帮助模型学习长距离范围内词元之间的关系。也可以把注意力看作一种把信息从一条残差流移动到另一条残差流的方法：它用另一个词元位置的信息来增强当前词元位置的流。

经过 $N$ 个 transformer 块之后，我们取最后一个 transformer 块产生的输出嵌入，让它通过一个线性反嵌入矩阵 $U$，再经过词汇表上的 softmax，从而生成所有可能下一个词元的概率分布。最后这两个组件，即反嵌入矩阵和 softmax，有时称为语言建模头。本章余下部分将更详细地介绍注意力以及这些模块。

图 8.2 展示了把 transformer 架构应用到上下文窗口 “So long and thanks for” 时的情形，并在每个词元位置显示最可能生成的下一个词元。在完整图中，$N$ 个块把整个上下文窗口的输入向量 $(x_1,\ldots,x_n)$ 映射为同样长度的输出向量窗口 $(h_1,\ldots,h_n)$。一列中可能包含从 12 到 96 个甚至更多堆叠的块。图中的箭头表示前面词元隐藏表示中的信息如何被纳入 transformer 块。

基于 transformer 的语言模型很复杂，因此细节会在本章和后续几章中逐步展开。第 7 章已经讨论过语言模型如何预训练，以及词元如何通过采样生成。本章余下部分将介绍多头注意力、transformer 块中的其他组件、输入编码，以及 transformer 的语言建模头。第 9 章介绍掩码语言建模和 BERT 系列双向 transformer 编码器模型。第 10 章说明如何对语言模型进行指令微调以执行 NLP 任务，以及如何让模型与人类偏好对齐。第 12 章会介绍带有编码器-解码器架构的机器翻译。第 15 章还会看到 transformer 在语音识别中的应用，以及编码器-解码器架构的进一步使用。

8.1   注意力

回忆第 5 章，对于 word2vec 和其他静态嵌入来说，一个词的意义表示总是同一个向量，与上下文无关。例如，单词 chicken 总是由同一个固定向量表示。因此，单词 it 的静态向量也许能编码出它是一个用于动物和无生命实体的代词。但在上下文中，it 的含义要丰富得多。考虑下面两句话中的 it：

(8.1) The chicken didn't cross the road because it was too tired.
(8.2) The chicken didn't cross the road because it was too wide.

在 (8.1) 中，it 指的是那只鸡，也就是说读者知道是鸡太累了；而在 (8.2) 中，it 指的是那条路，读者知道是路太宽了。^{[注 1]} 换言之，如果我们要计算这句话的意义，就需要让第一个句子中的 it 与 chicken 关联，让第二个句子中的 it 与 road 关联，并且这种关联必须对上下文敏感。

进一步说，设想像因果语言模型一样从左到右阅读，处理到单词 it 为止：

(8.3) The chicken didn't cross the road because it

在这一刻，我们还不知道 it 最终会指代什么。因此，此时 it 的表示可能同时具有 chicken 和 road 的某些方面，因为读者正在试图猜测接下来会发生什么。

词与相距很远的其他词之间存在丰富语言关系这一事实，贯穿于语言之中。再看两个例子：

(8.4) The keys to the cabinet are on the table.
(8.5) I walked along the pond, and noticed one of the trees along the bank.

在 (8.4) 中，短语 The keys 是句子的主语；在英语和许多语言中，主语必须在语法数上与动词 are 一致，此处二者都是复数。在英语中，不能用单数动词 is 搭配复数主语 keys（第 18 章会更详细讨论一致关系）。在 (8.5) 中，我们知道 bank 指的是池塘或河流的岸边，而不是金融机构，这是由上下文决定的，其中包括 pond 这样的词。（第 9 章会进一步讨论词义。）

这些例子的要点是：帮助我们在上下文中计算词义的上下文词，可能在句子或段落中距离目标词很远。Transformer 可以通过整合这些有用上下文词的含义，构建词义的上下文表示，即上下文嵌入。在 transformer 中，模型逐层建立越来越丰富的输入词元意义的上下文化表示。在每一层，我们会把来自上一层的词元 $i$ 的信息与相邻词元的信息结合起来，为每个位置上的每个词生成上下文化表示。

注意力就是 transformer 中的机制：它在从第 $k$ 层构建第 $k+1$ 层词元表示时，对上下文中合适的其他词元表示进行加权并组合。

图 8.3 展示了一个从 transformer 简化而来的示意例子。该图描述的是当前词元为 it 时，我们需要在 transformer 的第 $k+1$ 层为它计算上下文表示，并利用每个先前词元在第 $k$ 层的表示。图中用颜色表示上下文词上的注意力分布：词元 chicken 和 road 都有较高的注意力权重。这意味着当我们计算 it 的表示时，会主要利用 chicken 和 road 的表示。这对构建 it 的最终表示很有用，因为它最终会与 chicken 或 road 共指。

下面转向这个注意力分布如何表示以及如何计算。

8.1.1   更形式化地看注意力

如前所述，注意力计算是一种为 transformer 某一层中的某个词元计算向量表示的方法：它选择性地关注并整合前一层中先前词元的信息。注意力接收与位置 $i$ 的输入词元对应的输入表示 $x_i$，以及先前输入 $x_1,\ldots,x_{i-1}$ 构成的上下文窗口，然后产生输出 $a_i$。

在因果的、从左到右的语言模型中，上下文是所有先前的词。也就是说，在处理 $x_i$ 时，模型可以访问 $x_i$，也可以访问上下文窗口中所有先前词元的表示（上下文窗口包含数千个词元），但不能访问 $i$ 之后的词元。作为对比，第 9 章会把注意力推广到也能看见未来词的情形。

图 8.4 展示了完整因果自注意力层中的信息流；同样的注意力计算会在每个词元位置 $i$ 并行发生。因此，一个自注意力层把输入序列 $(x_1,\ldots,x_n)$ 映射为同样长度的输出序列 $(a_1,\ldots,a_n)$。

注意力的简化版本

从本质上讲，注意力只是上下文向量的加权和；复杂之处主要在于权重如何计算、哪些东西被求和。为了教学目的，我们先描述注意力的简化直觉：位置 $i$ 处的注意力输出 $a_i$ 就是所有表示 $x_j$（其中 $j\leq i$）的加权和；用 $\alpha_{ij}$ 表示 $x_j$ 应该对 $a_i$ 贡献多少：

每个 $\alpha_{ij}$ 都是一个标量，用来在对输入求和以计算 $a_i$ 时，对输入 $x_j$ 的值进行加权。那么如何计算这个权重 $\alpha$ 呢？在注意力中，我们会按照每个先前嵌入与当前词元 $i$ 的相似程度来加权。因此，注意力的输出，是先前词元嵌入的加权和；权重来自它们与当前词元嵌入的相似度。我们通过点积计算相似度分数；点积把两个向量映射为一个从 $-\infty$ 到 $\infty$ 的标量值。分数越大，被比较的向量越相似。然后用 softmax 对这些分数归一化，得到权重向量 $\alpha_{ij}, j\leq i$。

因此，在图 8.4 中，我们通过计算三个分数 $x_3\cdot x_1$、$x_3\cdot x_2$ 和 $x_3\cdot x_3$ 来得到 $a_3$，然后用 softmax 对它们归一化，并把得到的概率作为权重，表示每个先前位置与当前位置 $i$ 的相对相关程度。当然，$x_i$ 对自身非常相似，点积通常很高，因此 softmax 权重很可能在 $x_i$ 上最高。但其他上下文词也可能与 $i$ 相似，softmax 也会给它们分配一定权重。然后，我们把这些权重作为式 (8.6) 中的 $\alpha$ 值，计算加权和，即 $a_3$。

式 (8.6)--(8.8) 中的简化注意力展示了基于注意力计算 $a_i$ 的基本思路：把 $x_i$ 与先前向量比较，把这些分数归一化为概率分布，再用这个分布对先前向量求加权和。现在可以去掉这些简化了。

使用查询、键和值矩阵的单个注意力头

理解了注意力的简单直觉之后，我们引入真正的注意力头，也就是 transformer 中使用的注意力版本。（在 transformer 中，head 常用来指特定结构化层。）注意力头让我们能够明确区分每个输入嵌入在注意力过程中扮演的三种不同角色：

• 作为当前元素，与先前输入进行比较。我们把这种角色称为查询（query）。
• 作为先前输入，与当前元素比较以确定相似度权重。我们把这种角色称为键（key）。
• 最后，作为先前元素的值（value），被加权并求和以计算当前元素的输出。

为了捕获这三种不同角色，transformer 引入权重矩阵 $W_Q$、$W_K$ 和 $W_V$。这些权重会把每个输入向量 $x_i$ 投影到它作为查询、键或值时的表示：

有了这些投影，当我们计算当前元素 $x_i$ 与某个先前元素 $x_j$ 的相似度时，会使用当前元素查询向量 $q_i$ 与先前元素键向量 $k_j$ 的点积。此外，点积结果可能是任意大的正值或负值，对大值取指数会在训练中造成数值问题和梯度损失。为了避免这一点，我们用一个与嵌入大小有关的因子缩放点积，即除以查询和键向量维度 $d_k$ 的平方根。因此，简化的式 (8.7) 被替换为式 (8.11)。随后得到 $\alpha_{ij}$ 的 softmax 计算保持不变，但现在 $head_i$ 的输出计算基于值向量 $v$ 的加权和（式 (8.13)）。

下面给出从单个输入向量 $x_i$ 计算单个自注意力输出向量 $a_i$ 的最终方程。这一版本的注意力通过对先前元素的值求和来计算 $a_i$；每个值的权重由它的键与当前元素查询之间的相似度决定：

注意，我们还引入了另一个矩阵 $W_O$，它与注意力头输出相乘。这是为了重新调整注意力头的输出形状。注意力的输入 $x_i$ 和注意力的输出 $a_i$ 都具有相同维度 $[1\times d]$。我们常把 $d$ 称为模型维度；正如第 8.2 节将讨论的，每个 transformer 块的输出 $h_i$ 以及块内部的中间向量也都具有同样维度 $[1\times d]$。所有东西都保持相同维度，使 transformer 具有很好的模块性。

下面讨论形状。如何从输入的 $[1\times d]$ 到输出的 $[1\times d]$？看一下所有内部形状。查询和键向量有维度 $d_k$。查询向量和键向量都是 $[1\times d_k]$，因此可以计算点积 $q_i\cdot k_j$ 得到一个标量。值向量有单独的维度 $d_v$。变换矩阵 $W_Q$ 的形状为 $[d\times d_k]$，$W_K$ 为 $[d\times d_k]$，$W_V$ 为 $[d\times d_v]$。因此，式 (8.13) 中 $head_i$ 的输出形状是 $[1\times d_v]$。为了得到所需的输出形状 $[1\times d]$，需要调整注意力头输出的形状，所以 $W_O$ 的形状是 $[d_v\times d]$。在原始 transformer 工作中（Vaswani et al., 2017），$d=512$，$d_k$ 和 $d_v$ 都是 64。

多头注意力

式 (8.11)--(8.13) 描述的是单个注意力头。但实际上，transformer 使用多个注意力头。直觉是，每个头可能出于不同目的关注上下文：有的头可能专门表示上下文元素与当前词元之间的不同语言关系，有的头可能寻找上下文中特定类型的模式。

因此，在多头注意力中，我们有 $A$ 个独立注意力头，它们位于模型同一深度的并行层中，每个头都有自己的参数集合，从而允许该头建模输入之间关系的不同方面。因此，自注意力层中的每个头 $i$ 都有自己的查询、键和值矩阵：$W_Q^i$、$W_K^i$ 和 $W_V^i$。这些矩阵用于把输入投影为每个头各自的查询、键和值嵌入。

使用多个头时，模型维度 $d$ 仍用于输入和输出，查询和键嵌入的维度是 $d_k$，值嵌入的维度是 $d_v$（同样，在原始 transformer 论文中，$d_k=d_v=64$，$A=8$，$d=512$）。因此，对每个头 $i$，权重层 $W_Q^i$ 的形状为 $[d\times d_k]$，$W_K^i$ 的形状为 $[d\times d_k]$，$W_V^i$ 的形状为 $[d\times d_v]$。

下面是加入多头后的注意力方程；图 8.6 给出了直观示意。

注意式 (8.20) 中，$\mathrm{MultiHeadAttention}$ 是当前输入 $x_i$ 以及所有其他输入的函数。对于本章使用的因果或从左到右的注意力，其他输入只在左侧；但第 9 章还会看到另一种注意力，它也可以把右侧词元作为函数输入。我们会在式 (8.34) 中引入右侧上下文掩码时回到因果输入这一思想。

$A$ 个头中每个头的输出形状都是 $[1\times d_v]$，因此带有 $A$ 个头的多头层输出由 $A$ 个形状为 $[1\times d_v]$ 的向量构成。它们被拼接成一个维度为 $[1\times Ad_v]$ 的输出。然后，我们使用另一个线性投影 $W_O\in\mathbb{R}^{Ad_v\times d}$ 对其重新调整形状，得到在每个输入 $i$ 处具有正确输出形状 $[1\times d]$ 的多头注意力向量 $a_i$。

8.2   Transformer 块

自注意力计算位于transformer 块的核心。除了自注意力层之外，transformer 块还包含另外三类层：1）前馈层，2）残差连接，3）归一化层，口语中常称为 layer norm。

图 8.7 展示了一个 transformer 块，并勾勒出一种常见思路，称为残差流（Elhage et al., 2021）。在残差流视角下，我们把单个词元 $i$ 通过 transformer 块的处理，看成位置 $i$ 上一条由 $d$ 维表示组成的流。这条残差流从原始输入向量开始，各组件从残差流中读取输入，并把自己的输出加回这条流。

流底部的输入是一个词元的嵌入，维度为 $d$。这个初始嵌入会通过残差连接向上传递，并被 transformer 的其他组件逐步添加信息：已经介绍过的注意力层，以及即将介绍的前馈层。在注意力层和前馈层之前，会有一种称为层归一化的计算。

因此，初始向量先经过层归一化和注意力层，然后结果被加回流中，在这里是加回原始输入向量 $x_i$。随后，这个求和后的向量又经过另一个层归一化和前馈层，其输出再被加回残差。我们用 $h_i$ 表示词元 $i$ 经过该 transformer 块后的结果输出。

前馈层

前馈层是一个全连接两层网络，即一个隐藏层和两个权重矩阵，正如第 6 章所介绍的。各词元位置 $i$ 使用相同的权重，但不同层之间权重不同。通常会让前馈网络隐藏层维度 $d_{\mathrm{ff}}$ 大于模型维度 $d$。（例如，在原始 transformer 模型中，$d=512$，$d_{\mathrm{ff}}=2048$。）

层归一化

在 transformer 块中的两个阶段，我们会对向量进行归一化（Ba et al., 2016）。这个过程称为层归一化（layer norm，是 layer normalization 的简称），它是深度神经网络中许多归一化形式之一，用于通过把隐藏层的值保持在便于基于梯度训练的范围内，来改进训练表现。

层归一化可以看作统计学中 z-score 的一种变体，应用于隐藏层中的单个向量。也就是说，layer norm 这个术语有点容易误导；它不是应用于整个 transformer 层，而只是应用于单个词元的嵌入向量。因此，层归一化的输入是一个维度为 $d$ 的单个向量，输出是该向量归一化后的结果，维度仍为 $d$。层归一化的第一步，是在要归一化的向量各元素上计算均值 $\mu$ 和标准差 $\sigma$。给定维度为 $d$ 的嵌入向量 $x$，这些值计算如下：

有了这些值之后，对向量分量进行归一化：从每个分量中减去均值，再除以标准差。这个计算的结果是一个均值为 0、标准差为 1 的新向量。

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

合在一起

一个 transformer 块计算的函数，可以分解为每个组件计算的一条方程。这里用 $t$（形状为 $[1\times d]$）表示 transformer，并用上标区分块内部的各步计算：

注意，唯一把其他词元的信息（也就是其他残差流）作为输入的组件，是多头注意力；如式 (8.27) 所示，它会查看上下文中的所有相邻词元。不过，注意力的输出随后会被加到当前词元的嵌入流中。事实上，Elhage et al. (2021) 表明，可以把注意力头看作在字面意义上把信息从相邻词元的残差流移动到当前流中。因此，每个位置上的高维嵌入空间包含关于当前词元以及相邻词元的信息，尽管这些信息位于向量空间的不同子空间中。图 8.8 展示了这种移动的可视化。因此，我们把注意力函数称为该架构的词元混合组件，因为它把相邻词元流中的信息混合到当前流中。

关键的一点是，transformer 块的输入维度和输出维度相匹配，因此可以堆叠。输入到块中的每个词元向量 $x_i$ 的维度是 $d$，输出 $h_i$ 的维度也为 $d$。大型语言模型中的 transformer 会堆叠许多这样的块：从 12 层（用于 T5 或 GPT-3-small 语言模型）到 96 层（用于 GPT-3 large），再到更新模型中甚至更多层。后文还会回到堆叠问题。

式 (8.26) 及其后续方程只描述了单个 transformer 块，但残差流这个隐喻贯穿所有 transformer 层，从第一个 transformer 块一直到 12 层 transformer 中的第 12 个块。在较早的 transformer 块中，残差流表示当前词元。在最高的 transformer 块中，残差流通常表示下一个词元，因为模型最终被训练来预测下一个词元。

一旦堆叠许多块，还需要满足另一个要求：在最后一个、最高的 transformer 块的末端，对每个词元流最后的 $h_i$ 再运行一次额外的层归一化，位置正好在稍后会定义的语言模型头层之下。^{[注 2]}

8.3   使用单个矩阵 $X$ 并行化计算

前面对多头注意力和 transformer 块其余部分的描述，是从单个时间步 $i$、单条残差流中计算单个输出的角度展开的。但如前所述，为每个词元计算 $a_i$ 的注意力计算彼此独立；transformer 块中从输入 $x_i$ 计算 $h_i$ 的所有计算也同样彼此独立。这意味着我们可以很容易地并行化整个计算，利用高效的矩阵乘法例程。

做法是把输入序列中 $N$ 个词元的输入嵌入打包成一个大小为 $[N\times d]$ 的矩阵 $X$。$X$ 的每一行都是一个输入词元的嵌入。大型语言模型的 transformer 通常具有从 1K 到 32K 的输入长度 $N$；通过一些特殊长上下文机制等架构改变，也可以实现 128K 甚至数百万词元的更长上下文，这里不讨论。因此，对于普通 transformer，可以把 $X$ 看成有 1K 到 32K 行，每行维度为嵌入维度 $d$，也就是模型维度。

并行化注意力

先看单个注意力头，然后再转向多个头，并加入 transformer 块中的其余组件。对于一个头，我们把 $X$ 分别乘以查询、键和值矩阵：形状为 $[d\times d_k]$ 的 $W_Q$，形状为 $[d\times d_k]$ 的 $W_K$，以及形状为 $[d\times d_v]$ 的 $W_V$。这样得到矩阵 $Q$（形状 $[N\times d_k]$）、$K$（形状 $[N\times d_k]$）和 $V$（形状 $[N\times d_v]$），其中包含所有查询、键和值向量：

给定这些矩阵，我们可以通过一次矩阵乘法 $QK^\top$ 同时计算所有所需的查询-键比较。这个乘积的形状是 $N\times N$，如图 8.9 所示。

得到 $QK^\top$ 矩阵后，我们可以非常高效地缩放这些分数、做 softmax，然后把结果乘以 $V$，得到形状为 $N\times d$ 的矩阵：也就是输入中每个词元的向量嵌入表示。这样，整个 $N$ 个词元序列、单个头的完整自注意力步骤，就被化简为如下计算：

屏蔽未来

你可能注意到，上面的式 (8.33) 引入了一个 mask 函数。这是因为目前描述的自注意力计算存在一个问题：$QK^\top$ 的计算会为每个查询值和每个键值生成一个分数，包括查询之后的那些键。在语言建模设定中，这是不合适的：如果已经知道下一个词，猜测下一个词就太容易了。为了解决这个问题，矩阵上三角部分的元素被设置为 $-\infty$，softmax 会把它们变成 0，从而消除序列中后续词的信息。实践中，这是通过加入一个掩码矩阵 $M$ 完成的，其中对所有 $j>i$，$M_{ij}=-\infty$（即上三角部分），否则 $M_{ij}=0$。图 8.10 展示了得到的掩码 $QK^\top$ 矩阵。（第 9 章会看到，在需要未来词的任务中如何利用未来词。）

图 8.11 以矩阵形式示意了单个注意力头并行化后的所有计算。

图 8.9 和图 8.10 也清楚表明，注意力在输入长度上是二次复杂度的，因为在每一层都需要计算输入中每对词元之间的点积。这使得对非常长的文档（如整部小说）计算注意力代价高昂。尽管如此，现代大型语言模型仍然能够使用数千或数万词元的相当长上下文。

并行化多头注意力

在多头注意力中，和自注意力一样，输入与输出具有模型维度 $d$，键和查询嵌入具有维度 $d_k$，值嵌入具有维度 $d_v$（仍以原始 transformer 论文为例，$d_k=d_v=64$，$A=8$，$d=512$）。因此，对每个头 $c$，我们有形状为 $[d\times d_k]$ 的权重层 $W_Q^c$、形状为 $[d\times d_k]$ 的 $W_K^c$，以及形状为 $[d\times d_v]$ 的 $W_V^c$。这些矩阵与打包为 $X$ 的输入相乘，得到形状为 $[N\times d_k]$ 的 $Q$、形状为 $[N\times d_k]$ 的 $K$ 和形状为 $[N\times d_v]$ 的 $V$。$A$ 个头中每个头的输出形状是 $[N\times d_v]$，所以带有 $A$ 个头的多头层输出由 $A$ 个 $[N\times d_v]$ 矩阵构成。为了在后续处理中使用这些矩阵，它们被拼接为一个维度为 $[N\times Ad_v]$ 的单个输出。最后，使用形状为 $[Ad_v\times d]$ 的最终线性投影 $W_O$，把它重新调整到每个词元的原始输出维度。将拼接后的 $[N\times Ad_v]$ 矩阵输出乘以形状为 $[Ad_v\times d]$ 的 $W_O$，得到形状为 $[N\times d]$ 的自注意力输出 $A$。

使用并行输入矩阵 $X$ 合并全部计算

由一整层 $N$ 个 transformer 块并行计算的函数，其中每个块对应 $N$ 个输入词元之一，可以表示为：

注意，在式 (8.38) 中，$X$ 表示该层的输入，无论它来自哪里。对于第一层，如下一节所述，该输入就是此前一直称为 $X$ 的初始词嵌入加位置嵌入向量。但对于后续层 $k$，输入是前一层的输出 $H_{k-1}$。也可以把一个 transformer 层执行的计算分解开来，为每个组件写一条方程。这里用 $T$（形状为 $[N\times d]$）表示 transformer，并用上标区分块内部的各步计算；同样用 $X$ 表示来自前一层的块输入或初始嵌入：

这里写 $\mathrm{FFN}(T^3)$ 这样的记号时，是指同一个 FFN 并行应用于窗口中的 $N$ 个嵌入向量。同样，$N$ 个词元也会在 LayerNorm 中并行归一化。关键的是，transformer 块的输入维度和输出维度相匹配，因此可以堆叠。由于块输入中的每个词元 $x_i$ 都由维度为 $[1\times d]$ 的嵌入表示，因此输入 $X$ 和输出 $H$ 的形状都是 $[N\times d]$。

8.4   输入：词元嵌入和位置嵌入

下面讨论输入 $X$ 从哪里来。给定一个包含 $N$ 个词元的序列（$N$ 是以词元计的上下文长度），形状为 $[N\times d]$ 的矩阵 $X$ 为上下文中的每个词包含一个嵌入。Transformer 通过分别计算两种嵌入来做到这一点：输入词元嵌入和输入位置嵌入。

词元嵌入在第 6 章已经介绍，是维度为 $d$ 的向量，将作为输入词元的初始表示。（当向量沿着残差流通过 transformer 层向上传递时，这个嵌入表示会发生变化并扩展，纳入上下文，并根据我们正在构建的语言模型类型扮演不同角色。）初始嵌入集合存储在嵌入矩阵 $E$ 中，词汇表中的每个 $|V|$ 个词元对应一行。（提醒：这里的 $V$ 表示词元词汇表，与值向量中的 $V$ 无关。）因此，每个词都是一个 $d$ 维行向量，$E$ 的形状为 $[|V|\times d]$。

给定输入词元字符串如 “Thanks for all the”，我们首先把词元转换为词汇表索引（这些索引是在最初使用 BPE 或 SentencePiece 对输入进行词元化时创建的）。所以 thanks for all the 的表示可能是 $w=[5,4000,10532,2224]$。接着，我们使用索引从 $E$ 中选择相应行，即第 5 行、第 4000 行、第 10532 行、第 2224 行。

从嵌入矩阵中选择词元嵌入的另一种理解方式，是把词元表示成形状为 $[1\times |V|]$ 的独热向量，也就是词汇表中每个词对应一个维度。回忆一下，在独热向量中，除了一个元素以外所有元素都是 0；非零元素所在维度是该词在词汇表中的索引，其值为 1。因此，如果单词 “thanks” 在词汇表中的索引是 5，那么 $x_5=1$，且对所有 $i\neq 5$，$x_i=0$，如下所示：

用只有一个非零元素 $x_i=1$ 的独热向量相乘，实际上就是选出单词 $i$ 对应的行向量，从而得到单词 $i$ 的嵌入，如图 8.12 所示。

可以把这个思想扩展到整个词元序列：用一个独热向量矩阵表示整个序列，其中 transformer 上下文窗口中的每个 $N$ 个位置各有一个独热向量，如图 8.13 所示。

这些词元嵌入本身不依赖位置。为了表示序列中每个词元的位置，我们把这些词元嵌入与输入序列中各位置对应的位置嵌入结合起来。

这些位置嵌入从哪里来？最简单的方法称为绝对位置，即从随机初始化的嵌入开始，为直到某个最大长度的每个可能输入位置分配一个嵌入。例如，正如我们为单词 fish 有一个嵌入，也会为位置 3 有一个嵌入。和词嵌入一样，这些位置嵌入会在训练中与其他参数一起学习。我们可以把它们存储在形状为 $[N\times d]$ 的矩阵 $E_{\mathrm{pos}}$ 中。

为了生成捕获位置信息的输入嵌入，只需把每个输入的词嵌入与对应的位置嵌入相加。单个词元嵌入和位置嵌入的大小都是 $[1\times d]$，所以它们的和也是 $[1\times d]$。这个新嵌入会作为进一步处理的输入。图 8.14 展示了这一思想。

输入的最终表示，即矩阵 $X$，是一个 $[N\times d]$ 矩阵，其中第 $i$ 行是输入中第 $i$ 个词元的表示。它通过把 $E[id(i)]$，也就是出现在位置 $i$ 的词元 id 的嵌入，加到 $P[i]$，也就是位置 $i$ 的位置嵌入上得到。

简单位置嵌入方法的一个潜在问题是，训练集中输入开头位置会有大量例子，而越接近最大长度边界的位置例子会更少。这些后部位置的嵌入可能训练不足，在测试中泛化较差。另一种做法是选择一个静态函数，把整数输入映射为实值向量，并且能更好地处理任意长度序列。原始 transformer 工作使用了不同频率的正弦和余弦函数组合。正弦位置嵌入还可能有助于捕获位置之间固有关系，比如输入中位置 4 与位置 5 的关系比与位置 17 的关系更近。

更复杂的位置嵌入方法会进一步扩展这种捕获关系的思想，直接表示相对位置而不是绝对位置；这种方法通常在每一层的注意力机制中实现，而不是只在最初输入时加一次。

8.5   语言建模头

还必须介绍 transformer 的最后一个组件：语言建模头。这里的 head 指的是当我们把预训练 transformer 模型应用到各种任务时，添加到基本 transformer 架构顶部的额外神经线路。语言建模头就是完成语言建模所需的线路。

回忆一下，从第 3 章的简单 n-gram 模型到第 6 章和第 13 章的前馈与 RNN 语言模型，语言模型都是词预测器。给定词的上下文，它们为每个可能的下一个词分配概率。例如，如果前面的上下文是 “Thanks for all the”，而我们想知道下一个词是 “fish” 的可能性有多大，就会计算：

语言模型让我们能够为每个可能的下一个词分配这样的条件概率，从而得到整个词汇表上的分布。第 3 章的 n-gram 语言模型根据目标词与前面 $n-1$ 个词共同出现的计数来计算词的概率，因此上下文大小为 $n-1$。对于 transformer 语言模型，上下文大小就是 transformer 的上下文窗口大小；它可以相当大，例如大型模型中的 32K 词元（使用特殊长上下文架构时，数百万词的上下文也可能实现）。

语言建模头的任务，是取最后一个词元 $N$ 在最终 transformer 层的输出，并用它预测位置 $N+1$ 处即将出现的词。图 8.15 展示了如何完成这项任务：取最后一层最后一个词元的输出，即形状为 $[1\times d]$ 的 $d$ 维输出嵌入，并生成词上的概率分布，模型将从该分布中选择一个词来生成。

图 8.15 中的第一个模块是线性层，它的任务是从 $h_N^L$ 投影到logit 向量或分数向量。这里 $h_N^L$ 表示最终块 $L$ 在位置 $N$ 上的输出词元嵌入，因此形状为 $[1\times d]$。logit 向量会为词汇表 $V$ 中每个 $|V|$ 个可能词给出一个分数。因此，logit 向量 $u$ 的维度为 $[1\times |V|]$。

这个线性层可以学习，但更常见的做法是把这个矩阵与嵌入矩阵 $E$ 的转置绑定。回忆权重绑定：我们在模型中的两个不同矩阵上使用同一组权重。因此，在 transformer 的输入阶段，嵌入矩阵（形状 $[|V|\times d]$）用于把词汇表上的独热向量（形状 $[1\times |V|]$）映射为嵌入（形状 $[1\times d]$）。而在语言模型头中，嵌入矩阵的转置 $E^\top$（形状 $[d\times |V|]$）用于从嵌入（形状 $[1\times d]$）映射回词汇表上的向量（形状 $[1\times |V|]$）。在学习过程中，$E$ 会被优化到能够同时很好地完成这两种映射。因此，我们有时把转置 $E^\top$ 称为反嵌入层，因为它执行的是反向映射。

softmax 层把 logits $u$ 转换为词汇表上的概率 $y$。

这些概率可以用于为给定文本分配概率等任务。但最重要的用法是生成文本，也就是从这些概率 $y$ 中采样一个词。我们可以采样最高概率词（“贪心”解码），也可以使用第 7.4 节或第 8.6 节中的其他采样方法。

无论哪种情况，只要从概率向量 $y$ 中选定某个条目 $y_k$，就生成索引为 $k$ 的词。

图 8.16 展示了单个词元 $i$ 的完整堆叠架构。注意，每个 transformer 层的输入 $x_i^\ell$ 等于前一层输出 $h_i^{\ell-1}$。

在结束本节前补充一个术语说明：有时你会看到用于这种单向因果语言模型的 transformer 被称为仅解码器模型（decoder-only model）。这是因为这个模型大致构成第 12 章中会用于机器翻译的 transformer 编码器-解码器模型的一半。（容易混淆的是，transformer 最初被提出时采用的是编码器-解码器架构；后来才通过只使用原始架构中的解码器部分，定义了因果语言模型的标准范式。）

8.6   更多采样方法

下面介绍的采样方法都有一些参数，使生成过程能够在两个重要因素之间折中：质量与多样性。强调最高概率词的方法，往往会生成被人评价为更准确、更连贯、更符合事实的文本，但也更乏味、更重复。给中等概率词稍多权重的方法，往往更有创造性、更多样，但事实性更弱，也更可能不连贯或质量较低。

8.6.1   Top-k 采样

Top-k 采样是贪心解码的一个简单推广。我们不再选择单个最可能生成的词，而是先把分布截断到最可能的前 $k$ 个词，对它们重新归一化以得到合法概率分布，然后在这 $k$ 个词内部根据重新归一化后的概率随机采样。更形式化地说：

1. 预先选择词数 $k$。
2. 对词汇表 $V$ 中的每个词，使用语言模型计算该词在上下文条件下的似然 $p(w_t\mid w_{\lt t})$。
3. 按似然对词排序，并丢弃所有不属于前 $k$ 个最高概率词的词。
4. 对这 $k$ 个词的分数重新归一化，形成合法概率分布。
5. 根据概率，在剩下这 $k$ 个最可能词中随机采样一个词。

当 $k=1$ 时，top-k 采样与贪心解码完全相同。把 $k$ 设为大于 1 的数，会让模型有时选择不一定最高概率、但仍然足够可能的词；这种选择能生成更多样、同时质量仍足够高的文本。

8.6.2   Nucleus 采样或 top-p 采样

Top-k 采样的一个问题是 $k$ 是固定的，但不同上下文中的词概率分布形状不同。如果设置 $k=10$，有时前 10 个词非常可能，包含了大部分概率质量；但另一些时候，概率分布较平坦，前 10 个词只包含很小一部分概率质量。

另一种方法称为 top-p 采样或nucleus 采样（Holtzman et al., 2020），它保留的不是前 $k$ 个词，而是概率质量中最靠前的 $p$ 部分。目标仍然相同：截断分布以去除极不可能的词。但通过度量概率而不是词的数量，我们希望这种度量能在非常不同的上下文中更稳健，动态增加或减少候选词池。

给定分布 $P(w_t\mid w_{\lt t})$，我们按概率从高到低排序。于是 top-p 词汇表 $V^{(p)}$ 是满足下式的最小词集合：

8.7   训练

上一章已经描述了语言模型的训练过程。回忆一下，大型语言模型使用交叉熵损失训练，也称为负对数似然损失。在时间 $t$，交叉熵损失是模型分配给训练序列中下一个词的概率的负对数，即 $-\log p(w_{t+1})$。

图 8.17 展示了一般训练方法。在每一步，给定所有先前词，最终 transformer 层会在整个词汇表上产生一个输出分布。训练时，模型分配给正确词的概率用于计算序列中每个项的交叉熵损失。训练序列的损失是整个序列上的平均交叉熵损失。网络中的权重通过梯度下降调整，以最小化训练序列上的平均交叉熵损失。

对于 transformer，每个训练项都可以并行处理，因为序列中每个元素的输出都是单独计算的。

大型模型通常通过把完整上下文窗口填满文本来训练（例如 GPT-4 为 4096 个词元，Llama 3 为 8192 个词元）。如果文档短于窗口长度，则把多个文档打包进一个窗口，并在它们之间放置特殊的文本结束词元。梯度下降的 batch size 通常非常大（最大的 GPT-3 模型使用 320 万词元的 batch size）。

8.8   处理规模问题

大型语言模型确实很大。例如，Meta 的 Llama 3.1 405B Instruct 模型有 4050 亿参数（它有 $L=126$ 层，模型维度 $d=16{,}384$，以及 $A=128$ 个注意力头），使用 15.6 TB 的文本词元训练，词汇表规模为 128K 个词元（Llama Team, 2024）。因此，关于如何理解 LLM 的扩展，尤其是在资源有限时如何实现它们，有大量研究。接下来的几节讨论如何思考规模（即缩放律的概念），以及让语言模型高效工作的若干重要技术，例如 KV cache 和参数高效微调（PEFT）。

8.8.1   缩放律

大型语言模型的性能主要由 3 个因素决定：模型大小（不计嵌入的参数数量）、数据集大小（训练数据量），以及用于训练的计算量。也就是说，我们可以通过增加参数（增加层数或扩大上下文，或二者兼具）、使用更多数据训练，或训练更多迭代来改进模型。

这些因素与性能之间的关系称为缩放律。粗略地说，大型语言模型的性能（损失）会随模型训练的这三个属性分别呈幂律变化。

例如，Kaplan et al. (2020) 发现，在参数、数据集或计算预算受限训练模型时，如果每种情况下另外两个属性保持不变，损失 $L$ 与非嵌入参数数量 $N$、数据集大小 $D$、计算预算 $C$ 之间有如下三个关系：

非嵌入参数数量 $N$ 可以大致按如下方式计算（忽略偏置，令 $d$ 为模型输入和输出维度，$d_{\mathrm{attn}}$ 为自注意力层大小，$d_{\mathrm{ff}}$ 为前馈层大小）：

因此，GPT-3 有 $n=96$ 层、维度 $d=12288$，参数数量约为 $12\times 96\times 12288^2\approx 1750$ 亿。

$N_c$、$D_c$、$C_c$、$\alpha_N$、$\alpha_D$ 和 $\alpha_C$ 的取值依赖于具体 transformer 架构、词元化方式和词汇表大小。因此，缩放律关注的通常不是所有精确数值，而是它们与损失之间的关系。^{[注 3]}

缩放律有助于决定如何把模型训练到特定性能。例如，可以观察训练曲线早期阶段，或在较小数据量上的性能，来预测如果增加更多数据或增大模型规模时损失会如何变化。缩放律的其他方面还可以告诉我们，在扩大模型时需要增加多少数据。

8.8.2   KV Cache

我们在图 8.11 和式 (8.34) 中看到过，注意力向量在训练时可以通过两次矩阵乘法非常高效地并行计算。下面重复这一形式：

遗憾的是，在推理时不能完全像训练时一样进行这种高效计算。这是因为推理时，我们会一次一个地迭代生成下一个词元。对于刚生成的新词元，记为 $x_i$，我们需要把它分别乘以 $W_Q$、$W_K$ 和 $W_V$，以计算它的查询、键和值。但重新计算所有先前词元 $x_{\lt i}$ 的键和值向量会浪费计算时间；在先前步骤中，我们已经计算过这些键和值向量。因此，不必重新计算它们，我们会在每次计算键和值向量时把它们存入内存中的 KV cache，之后需要时直接从缓存中取出即可。图 8.18 修改了图 8.11，以展示单个新词元的计算过程，并标出哪些值可以从缓存中取出而不是重新计算。

8.8.3   参数高效微调

如上所述，一种很常见的做法是拿一个语言模型，在某个新领域的附加数据上微调它，即继续训练它来预测后续词，以便给模型更多关于该领域的信息。

微调非常大的语言模型可能很困难，因为需要训练的参数数量巨大；每一次 batch 梯度下降都要通过许多巨大的层进行反向传播。这使得微调巨型语言模型在算力、内存和时间上都极其昂贵。

因此，存在一些替代方法，使模型能够在不改变所有参数的情况下进行微调。这样的方法称为参数高效微调，有时称为 PEFT，因为我们会高效选择微调时要更新的参数子集。例如，可以冻结某些参数（不改变它们），只更新某个特定参数子集。

这里介绍其中一种模型，称为 LoRA，即低秩适配（Low-Rank Adaptation）。LoRA 的直觉是，transformer 有许多执行矩阵乘法的稠密层，例如注意力计算中的 $W_Q$、$W_K$、$W_V$ 和 $W_O$ 层。使用 LoRA 时，微调过程中不更新这些层，而是冻结它们，并改为更新一个参数更少的低秩近似。

考虑一个维度为 $[k\times d]$ 的矩阵 $W$，它原本需要在微调中通过梯度下降更新。通常，这个矩阵会得到维度为 $[k\times d]$ 的更新 $\Delta W$，用于在梯度下降之后更新 $k\times d$ 个参数。在 LoRA 中，我们冻结 $W$，转而更新 $W$ 的低秩分解。创建两个矩阵 $A$ 和 $B$，其中 $A$ 的大小是 $[k\times r]$，$B$ 的大小是 $[r\times d]$，并选择很小的 $r$，使得 $r\ll \min(d,k)$。微调时，我们更新 $A$ 和 $B$，而不是 $W$。也就是说，用 $W+AB$ 替代 $W+\Delta W$。图 8.19 展示了这一直觉。原本前向传播为 $h=xW$，新的前向传播改为：

LoRA 有许多优点。它显著降低硬件需求，因为大多数参数不需要计算梯度。权重更新可以直接加到预训练权重中，因为 $AB$ 与 $W$ 大小相同。这意味着它不会增加推理时间。它还意味着可以为不同领域构建 LoRA 模块，并通过把它们加到 $W$ 中或从 $W$ 中减去，来方便地切换。

在最初版本中，LoRA 只应用于注意力计算中的矩阵（$W_Q$、$W_K$、$W_V$ 和 $W_O$ 层）。现在存在许多 LoRA 变体。

8.9   解释 Transformer

基于 transformer 的语言模型究竟如何在语言任务中表现得如此出色？可解释性这个子领域，有时称为机制可解释性，关注如何从机制层面理解 transformer 内部发生了什么。接下来的两个小节讨论 transformer 可解释性的两个研究较多的方面。

8.9.1   上下文学习与归纳头

作为让模型做我们想做之事的一种方式，提示可以被看作与预训练根本不同。通过预训练学习，意味着根据某个损失函数，使用梯度下降更新模型参数。但带有示例的提示也能教模型完成新任务。模型在处理提示时，会从这些示例中学习关于任务的东西。

即使没有示例，也可以把提示过程看作一种学习。例如，模型在提示中走得越远，通常越擅长预测接下来出现的词元。上下文中的信息正在帮助模型获得更强预测能力。

上下文学习这个术语最早由 Brown et al. (2020) 在介绍 GPT-3 系统时提出，用于指语言模型从提示中进行的这两类学习。上下文学习指语言模型在推理时的前向传播过程中学习完成新任务、更好预测词元，或总体上降低损失，而不对模型参数进行任何基于梯度的更新。

上下文学习如何起作用？我们还不确定，但有一些很有启发性的想法。其中一个假设基于归纳头（Elhage et al., 2021; Olsson et al., 2022）的思想。归纳头是某种线路的名称，而线路是网络中的一种抽象组件。归纳头线路是 transformer 注意力计算的一部分，它最初是通过观察只有 1--2 个注意力头的小型语言模型发现的。

归纳头的功能是预测重复序列。例如，如果它在输入序列中看到模式 $AB\ldots A$，就会预测接下来是 $B$，从而实例化模式补全规则 $AB\ldots A\rightarrow B$。它通过注意力计算中的前缀匹配组件做到这一点：当看当前词元 $A$ 时，回溯上下文寻找先前出现过的 $A$。如果找到，归纳头会有一个复制机制，通过提高 $B$ 接下来出现的概率，“复制”那个较早 $A$ 后面的词元 $B$。图 8.20 展示了一个例子。

Olsson et al. (2022) 提出，这种模式补全规则的模糊泛化版本，例如实现类似 $A^*B^*\ldots A\rightarrow B$ 的规则，其中 $A^*\approx A$ 且 $B^*\approx B$（这里 $\approx$ 表示它们在某种意义上语义相似），可能负责上下文学习。支持这一假设的证据来自 Crosbie and Shutova (2022)：他们表明，消融归纳头会导致上下文学习性能下降。Ablation 原本是医学术语，意为移除某物。在 NLP 可解释性研究中，我们把它用作测试因果效应的工具：如果去掉一个假定的原因，就预期相应效果会消失。Crosbie and Shutova (2022) 通过先找到在随机输入序列上表现为归纳头的注意力头，然后把这些头的输出矩阵中的某些项设为零，从而把这些头的输出置零。结果确实发现，消融后的模型在上下文学习上差得多：它们从提示中的示例学习的能力显著下降。

8.9.2   Logit Lens

另一个有用的可解释性工具是 logit lens（Nostalgebraist, 2020）。它提供了一种可视化 transformer 内部层可能正在表示什么的方法。

其思想是：取 transformer 任意一层中的任意向量，假装它就是最终嵌入，然后直接把它乘以反嵌入层以得到 logits，再计算 softmax，查看该向量可能表示的词分布。这样可以成为观察模型内部表示的一扇有用窗口。由于网络并没有被训练成让内部表示以这种方式工作，logit lens 并不总是完美有效，但它仍然是帮助我们可视化 transformer 内部层的一个有用技巧。

8.10   小结

本章介绍了 transformer 及其用于上一章语言建模任务的各个组件。下面总结本章覆盖的主要要点：

• Transformer 是基于多头注意力的非循环网络；多头注意力是一种自注意力。多头注意力计算接收输入向量 $x_i$，并通过加入来自先前词元的向量把它映射为输出 $a_i$；这些先前词元向量的权重由它们与当前词处理的相关程度决定。
• 一个 transformer 块由残差流构成：来自前一层的输入向上传到下一层，不同组件的输出被加到这条流上。这些组件包括一个多头注意力层和随后的前馈层，并且二者前面各有层归一化。Transformer 块可以堆叠，从而形成更深、更强大的网络。
• Transformer 的输入通过把嵌入（由嵌入矩阵计算）与位置编码相加得到，位置编码表示词元在窗口中的序列位置。
• 可以由堆叠的 transformer 块构建语言模型，并在顶部放置语言模型头。语言模型头对最高层输出 $H$ 应用反嵌入矩阵以生成 logits，然后通过 softmax 生成词概率。
• 基于 transformer 的语言模型具有较宽的上下文窗口。对于带有特殊机制的超大模型，上下文窗口可达 200K 个词元甚至更多，使模型能够利用大量上下文来预测即将出现的词。
• 有多种计算技巧可以让大型语言模型更高效，例如 KV cache 和参数高效微调。

历史注记

Transformer（Vaswani et al., 2017）的发展借鉴了两条早期研究线索：自注意力和记忆网络。

编码器-解码器注意力，即使用输入词编码上的软权重来为生成式解码器提供信息的思想（见第 12 章），最早由 Graves (2013) 在手写生成语境中发展，并由 Bahdanau et al. (2015) 用于机器翻译。通过取消单独编码和解码序列的需求，并把注意力看作在把信息从低层传到高层时对词元进行加权的方式，这一思想被扩展为自注意力（Ling et al., 2015; Cheng et al., 2016; Liu et al., 2016）。

Transformer 的其他方面，包括 key、query 和 value 这些术语，来自记忆网络。记忆网络是一种向网络添加外部读写记忆的机制，它使用查询嵌入来匹配表示联想记忆中内容的键（Sukhbaatar et al., 2015; Weston et al., 2015; Graves et al., 2014）。

更多历史内容将在后续草稿中补充。