Transformer 架构

• 80 min read • 15909 words
Tags: Deep Learning NLP Transformer
Categories: NLP

Transformer 架构

1. 最小化的自注意力架构

a.a. 注意力的广义定义

我们可以将注意力机制（Attention）理解为一个过程，它模仿了我们从一个“键值对（Key-Value）”存储中“软性地”查找信息的方式：

  1. 我们有一个查询（Query）。
  2. 我们用这个查询去和存储中所有的键（Key）进行比较，找出哪些键和查询最相似。
  3. 我们并非只选择最相似的那一个，而是“软性地挑选”，即对所有的值（Value）进行加权平均。
  4. 权重的大小取决于对应的键和查询的相似度。键越相似，对应的值所占的权重就越大。

自注意力是注意力机制的一种特殊形式：用于定义查询、键和值的元素都来自于输入序列本身。换句话说，这是一个序列自己对自己进行注意力计算的过程。

b.b. 键-查询-值自注意力机制

对于输入序列中的每一个词向量 xjx_j，我们通过三个不同的、可学习的权重矩阵 QQ、KK、VV，将其映射成三个不同的向量：

  1. 查询向量：qi=Qxiq_i = Q x_i
    • qiq_i 代表 xix_i 为了更好地理解自己而去主动发出的查询。
  2. 键向量：kj=Kxjk_j = K x_j
    • kjk_j 代表 xjx_j 所携带信息的索引，用于响应查询。
  3. 值向量：vj=Vxjv_j = V x_j
    • vjv_j 代表 xjx_j 实际包含的、将要被提取的内容或信息。

通过使用三个不同的矩阵，模型可以学习将同一个输入向量 xix_i 投射到三个不同的语义空间中，让它在扮演“提问者”、“被检索的标签”和“信息提供者”这三个不同角色时，可以有不同的侧重和表示。

为了计算 xix_i 的新表示，我们需要决定序列中其他每个词 xjx_j 应该贡献多大的“注意力”。这个权重 αijα_{ij} 通过两步计算得出：

  1. 计算相似度：将 xix_i 的查询向量 qiq_i 与序列中每一个词 xjx_j 的键向量 kjk_j 进行点积运算 qiTkjq_i^T k_j。

这个点积结果衡量了 qiq_i 和 kjk_j 之间的相似度，也称为注意力得分。

  1. 归一化：将上一步得到的所有注意力得分 qiTk1,qiTk2,...,qiTknq_i^T k_1, q_i^T k_2, ..., q_i^T k_n作为一个整体，送入一个 softmax 函数。softmax 会将这些得分转换成一个概率分布，即最终的注意力权重 αijα_{ij}。所有权重加起来等于 1。

然后，xix_i 的最终上下文表示 hih_i 就是对序列中所有词的值向量 vjv_j 进行的加权求和。而权重就是我们刚刚算出的注意力权重 αijα_{ij}：

hi=∑j=1nαijvjh_i = \sum _{j=1}^n\alpha_{ij}v_j

这个操作会对序列中的每一个词 xix_i 都执行一遍，从而得到整个序列的上下文表示 h1:nh_{1:n}。最关键的是，这个过程摆脱了 RNN 的顺序依赖，可以高度并行化计算。

c.c. 位置表示

我们在上一节中提出的自注意力机制有下面的问题：它无法感知词的顺序（这可以由 hih_i 的计算过程看出）。

为了让自注意力模型理解顺序，我们可以通过下面两种方法将位置信息注入模型：

  1. 在输入中使用携带位置信息的向量。
  2. 修改自注意力操作本身，使其能够感知位置。
i.i. 位置嵌入

位置嵌入是最经典、最常用的方法，BERT 和原始 Transformer 都采用此方案的变体。

为了实现位置嵌入，我们创建一个新的、可学习的位置嵌入矩阵 PP。PP 的大小为 N×dN × d，其中 NN 是模型能处理的最大序列长度，dd 是嵌入维度。PP 的第 ii 行 PiP_i 就是一个代表“位置 ii”的向量。

这样，我们只需将一个词的词嵌入 xix_i 和它所在位置的位置嵌入 PiP_i 直接相加：

x^i=xi+Pi\hat{x}_i = x_i + P_i

经过这个操作后，新的输入向量 x^i\hat{x}_i 就同时包含了词义信息和位置信息。

iiii 修改注意力权重 α\alpha

修改注意力权重 α\alpha 是一种更现代的方法，它不去修改输入，而是直接在注意力计算过程中加入位置的“偏置”。

这个方法基于下面的核心直觉：在其他条件都相同的情况下，模型应该更关注“邻近”的词，而不是“遥远”的词。这是一种很强的归纳偏置（inductive bias）。

这一方法的一个简单的实现方式如下：在计算出原始的注意力得分（k_1:n * q_i）之后，但在送入 softmax 之前，我们给这些得分加上一个代表相对距离的偏置向量：

αi=softmax(k1:nqi+[−i,…,−1,0,−1,…,−(n−i)])\alpha_i = \text{softmax} \left( k_{1:n}q_i + [-i, \dots, -1, 0, -1, \dots, -(n - i)] \right)

这个偏置向量以当前词 ii 为中心（偏置为0），离得越远的词，其偏置值就越负，从而拉低了它的注意力得分。这相当于直接告诉模型：“优先关注你旁边的词”。

d.d. 矩阵形式的自注意力

到目前为止，我们都是以单个向量 hih_i 的视角来描述计算过程。在实践中，为了利用 GPU 的并行计算能力，我们可以将整个序列的计算过程写成矩阵运算的形式：

我们可以将之前对单个向量的操作，扩展为对整个矩阵 XX 的操作，从而一次性得到所有词的查询、键和值矩阵 QmatQ_{mat}, KmatK_{mat}, VmatV_{mat}：

Qmat=XWQQ_{mat} = X W_Q Kmat=XWKK_{mat} = X W_K Vmat=XWVV_{mat} = X W_V

这里的 WQW_Q, WKW_K, WVW_V 就是维度为 d×dd × d 的可学习权重矩阵。QmatQ_{mat}, KmatK_{mat}, VmatV_{mat} 三个矩阵的维度也都是 n×dn × d。

同时，我们也可以通过一次矩阵乘法计算出所有查询与所有键之间的相似度得分：

Scores=QmatKmatT\text{Scores} = Q_{mat} K_{mat}^T

QmatQ_{mat} 维度为 n×dn × d，KmatTK_{mat}^T 的维度为 d×nd × n，这样我们就得到一个 n×nn × n 的得分矩阵 ScoresScores，其中 Scores[i,j]\text{Scores}[i, j] 对应 qiq_i 和 kjk_j 之间的原始得分。

之后我们对得分矩阵按行进行 softmax 操作，得到注意力权重矩阵 AA：

A=softmax(ScaledScores)A = \text{softmax}(Scaled_{Scores})

AA 的维度为 n×nn × n，其中A[i,j]A[i, j] 代表词 ii 对词 jj 的注意力权重。

我们将权重矩阵 AA 与值矩阵 VmatV_{mat} 相乘，就可以得到最终的上下文表示矩阵 HH：

H=A∗VmatH = A * V_mat

VmatV_{mat} 的维度为 n×dn × d，得到的最终输出矩阵 HH 维度为n×dn × d。HH 的第 ii 行就是我们之前求的 hih_i。

e.e. 按元素的非线性变换

i.i. 非线性变换

如果我们仅仅是将自注意力层一个接一个地堆叠起来，会缺失深度学习架构中一个至关重要的部分：按元素的非线性变换（elementwise nonlinearities）。

自注意力机制的核心操作本质上是线性的。虽然 softmax 引入了非线性，但它作用于权重，而对 ValueValue 向量的组合是线性的。假设我们有两层自注意力。第二层的输入是第一层的输出，通过下面的计算，最终可以被重新组合成一个等效的、单层的自注意力形式。

oi=∑j=1nαijV(2)(∑k=1nαjkV(1)xk)=∑k=1n(∑j=1nαijαjk)V(2)V(1)xk=∑k=1nαik∗V∗xk\begin{aligned} o_i &= \sum_{j=1}^{n} \alpha_{ij} V^{(2)} \left( \sum_{k=1}^{n} \alpha_{jk} V^{(1)} x_k \right) \\ &= \sum_{k=1}^{n} \left( \sum_{j=1}^{n} \alpha_{ij} \alpha_{jk} \right) V^{(2)} V^{(1)} x_k \\ &= \sum_{k=1}^{n} \alpha^{*}_{ik} V^{*} x_k \end{aligned}

可以看到，线性变换的堆叠，其结果仍然是一个线性变换。。

为了解决这个问题，在实践中，标准的做法是在每个自注意力层之后，应用一个前馈网络中：

hFF=W2∗ReLU(W1hself−attention+b1)+b2h_{FF} = W^2 * \text{ReLU}(W^1 h_{self-attention} + b_1) + b_2

正是 ReLU 激活函数引入了强大的非线性，打破了前面自注意力层的线性局限，极大地提升了注意力机制的表达能力。

ii.ii. 升维与降维

在非线性变换的前馈网络中，中间隐藏层的维度通常会远大于模型的主维度 dd。例如：我们可以使用下面两个矩阵：W1W_1 将 dd 维的输入映射到一个更高的维度（比如 4d4d 或 5d5d），然后 W2W_2 再将其映射回 dd 维。

这使得模型能够在一个更高维的“隐藏空间”中学习到更丰富、更复杂的特征，然后再将这些特征压缩回原始维度，传递给下一层。

f.f. 未来信息掩码

未来信息掩码主要用于自回归建模 (Autoregressive Modeling) 问题：根据到目前为止已经出现的所有词 w1:t−1w_{1:t-1}，来预测下一个词 wtw_t：

wt∼softmax(f(w1:t−1))w_t \sim \text{softmax}(f(w_{1:t-1}))

这个过程中最至关重要的一点是，在预测未来时，绝对不能偷看未来的信息。否则，这个问题就变得毫无意义了

在传统的循环神经网络（RNN）中，这个“不偷看未来”的规则是内建于其结构中的：当我们要预测第 tt 个词 wtw_t 时，我们使用的是第 t−1t-1 个时间步的隐藏状态 ht−1h_{t-1}。 。但是在 Transformer 的自注意力机制中，默认情况下，自注意力是“全局”的，每个词都会关注序列中的所有词。

为了在 Transformer 中强制执行“不偷看未来”的规则，我们采用了一种简单而有效的方法：在计算注意力权重时，直接对未来的位置进行掩码（mask）：

  • 在将注意力得分送入 softmax 函数之前，我们给所有代表未来的位置（即 j>ij > i 的位置）的得分，加上一个非常大的负常数。

softmax 函数的计算涉及指数 exp⁡()\exp()。一个非常大的负数的指数会无限趋近于 0。因此，经过 softmax 之后，这些未来位置的注意力权重 αijα_{ij} 就会变为 0。

f.f. 自注意力架构总结

综上，我们可以得到：一个可用的、最小化的自注意力架构，包含以下四个关键组成部分：

  1. 自注意力操作：这是整个架构的核心，通过“查询-键-值”模型，让序列中的每个词都能直接与其他所有词进行交互，从而构建上下文表示。
  2. 位置表示：通过引入位置嵌入等技术，向模型注入了关于词序的关键信息，使其能够区分“面包烤了烤箱”和“烤箱烤了面包”。
  3. 按元素的非线性变换：在架构中加入非线性变换（例如使用 ReLU 激活函数的全连接前馈网络），来提升模型的复杂度和表达能力。
  4. 未来信息掩码：这个组件是有条件的、针对特定任务的。它专门用于自回归语言模型的场景，比如文本生成。

2. Transformer

Transformer 是一个基于自注意力的架构，它由多个堆叠的模块组成。

a.a. 多头自注意力

i.i. 基本概念

一个单独的自注意力层，在对一个词进行信息汇总时，倾向于只关注一种“平均”的相似性关系。如果想同时关注多种不同类型的关系（比如，一个词的句法依赖关系和它的语义相似关系），单头注意力会很难做到，因为它需要在计算点积时需要权衡这几种不同类型的关系。

而多头注意力可以很好地解决这个问题：与其让一个注意力机制“分心”去做多件事情，不如设置多个独立的注意力头，让每个头都使用不同的查询、键、值转换矩阵，去关注输入信息中不同的方面或子空间。最后，再将所有头的结果综合起来。

ii.ii. 计算过程

假设我们有 kk 个注意力头，我们执行如下的操作：

  1. 为每个头定义独立的变换矩阵：对于第 ℓ\ell 个头（ℓ\ell 从 1 到 kk），我们都有一组独立的权重矩阵 Q(ℓ),K(ℓ),V(ℓ)Q(\ell), K(\ell), V(\ell)。
  2. 降维：一个关键的设计是，这些矩阵会将原始的 dd 维输入，投影到一个更低的维度 dk\frac{d}{k}。
  3. 并行计算注意力：让每个头都独立地执行一次完整的自注意力计算。第 ℓ\ell 个头的输出 h(ℓ)ih(\ell)_i 是一个 dk\frac{d}{k} 维的向量：
hi(ℓ)=∑j=1nαij(ℓ)vj(ℓ)h_i^{(\ell)} = \sum_{j=1}^{n} \alpha_{ij}^{(\ell)} v_j^{(\ell)} αij(ℓ)=exp⁡(qi(ℓ)⊤kj(ℓ))∑j′=1nexp⁡(qi(ℓ)⊤kj′(ℓ))\alpha_{ij}^{(\ell)} = \frac{\exp(q_i^{(\ell)\top} k_j^{(\ell)})}{\sum_{j'=1}^{n} \exp(q_i^{(\ell)\top} k_{j'}^{(\ell)})}
  1. 拼接与整合：
    • 将所有 kk 个头的输出向量 hi(1),...,hi(k)h^{(1)}_i, ..., h^{(k)}_i 拼接在一起。因为每个向量都是 dk\frac{d}{k} 维，拼接后会得到一个完整的 dd 维向量。
    • 最后，再用一个额外的、可学习的线性权重矩阵 OO （维度为 d×dd \times d）对这个拼接后的向量做一次最终的变换，得到多头注意力的最终输出 hih_i：
hi=Oh(vi(1);… ;vi(k))h_i = O_h(v_i^{(1)}; \dots; v_i^{(k)})
iii.iii. 具体实现

在单头注意力中，我们通过 XQXQ, XKXK, XVXV 得到 n×dn \times d 维的矩阵。

在多头注意力中，我们不需要创建 kk 组小的矩阵并计算 kk 次。我们仍然先计算出大的 n×dn×d 维的 QmatQ_{mat}, KmatK_{mat}, VmatV_{mat} 矩阵。

这时我们需要执行一个关键步骤：重塑(Reshape)。我们将这个 n×dn×d 的矩阵，重塑成一个 n×k×dkn × k × \frac{d}{k} 的三维矩阵。这在数学上等价于将 dd 维的特征拆分成了 kk 组，每组 dk\frac{d}{k} 维。

通过这种方式，我们可以利用深度学习框架中的批处理运算，将 kk 个头的计算完全并行化，所有的矩阵乘法和 softmax 都是在 kk 个头上同时发生的。

b.b. 层归一化

i.i. 基本概念

在一个深度神经网络中，每一层的输出（激活值）会作为下一层的输入。如果某一层输出的数值范围变化很大，就会给下一层的学习带来困难，这个现象有时被称为“内部协方差偏移 (Internal Covariate Shift)”。

层归一化的直觉就是通过规范化每一层的激活值，减少其中无用的、剧烈的变动，从而为下一层提供一个更稳定、更容易学习的输入。

层归一化最大的好处可能并不仅仅是稳定了前向传播中的数据，而是在反向传播过程中改善了梯度。一个稳定、规范化的数据范围可以使得梯度更加平滑，不容易出现梯度消失或梯度爆炸的问题，从而让整个模型的训练更加高效和可靠。

ii.ii. 工作流程

层归一化的具体流程如下：

  1. 计算统计数据: 对某一层的所有激活值，计算它们的均值和方差。这是归一化的基础：
μ^i=1d∑j=1dhij,σ^i=1d∑j=1d(hij−μi)2\hat{\mu}_i = \frac{1}{d} \sum_{j=1}^{d} h_{ij}, \quad \hat{\sigma}_i = \sqrt{\frac{1}{d} \sum_{j=1}^{d} (h_{ij} - \mu_i)^2}
  • μ^i\hat{\mu}_i 是第 ii 个 token 的 所有 dd 个维度的平均值。注意，μ^i\hat{\mu}_i 是一个标量，σ^i\hat{\sigma}_i 同理。
  1. 执行归一化: 使用上一步计算出的均值和方差，对该层的激活值进行标准化处理，使得处理后的数据均值为0，方差为1：
LN(hi)=hi−μ^iσ^i\text{LN}(h_i) = \frac{h_i - \hat{\mu}_i}{\hat{\sigma}_i}
  • hih_i 代表第 ii 个 token 的完整向量（一个 dd 维向量）。

这个过程在数学上称为广播 (broadcasting)，即将标量 μ^i\hat{\mu}_i 和 σ^i\hat{\sigma}_i 扩展到与向量 hih_i 相同的维度进行计算。

iii.iii. 对层归一化的深入理解

在 Transformer 中，层归一化是独立地对序列中的每一个位置 (token) 进行的。统计数据是在该 token 的所有隐藏维度上计算的。

换言之，每个词的层归一化之间是完全独立的。统计是在特征维度上进行的，而不是在序列长度上。

c.c. 残差连接

残差连接的思路非常简单：将一个模块的输入，直接加到这个模块的输出上：

fresidual(h1:n)=f(h1:n)+h1:nf_{residual}(h_{1:n}) = f(h_{1:n}) + h_{1:n}

这就像是在一个复杂的处理流程旁边，开辟了一条“直通车道”，让原始信息可以直接流向下一层。

残差连接操作基于如下的思想：

  • 优化梯度流: 在非常深的网络中，梯度在反向传播时需要穿过很多层，每穿过一层就可能因为矩阵乘法等操作而变小，导致最开始的几层网络几乎学不到东西（这就是梯度消失问题）。而残差连接创造的这条“直通车道” +h+h 部分，其数学本质是一个恒等函数，对输入的导数（局部梯度）恒等于1。这意味着，梯度可以通过这条“直通车道”毫无衰减地向后传播。这极大地缓解了梯度消失问题。
  • 让学习变得更容易: 假设在某个场景下，最优的变换就是什么都不做（f(h)=hf(h) = h）。对于一个没有残差连接的普通网络，它需要费很大力气学习一个复杂的恒等变换。但有了残差连接，网络只需要学习让 f(h)f(h) 的输出接近于0就可以了。让一个网络输出0，远比让它学习一个完整的恒等变换要容易得多。
    • 推广开来，网络不需要从头学习一个完整的函数，而只需要学习原始输入与目标输出之间的“差异”或“残差”，这会大大降低学习难度。

d.d. Add&Norm 模块

Transformer 中的 Add&Norm 模块将层归一化和残差连接组合在一起。它们有下面两种组合方式：

  • 前置归一化：先对输入 hh 进行层归一化，然后将结果送入主模块 ff，最后再与原始输入 hh 进行残差连接：
hpre-norm=f(LN(h))+hh_{\text{pre-norm}} = f(\text{LN}(h)) + h
  • 后置归一化：先将输入 hh 送入主模块 ff，然后与原始输入 hh 进行残差连接，最后对相加的结果进行层归一化：
hpost-norm=LN(f(h)+h)h_{\text{post-norm}} = \text{LN}(f(h) + h)

研究发现前置归一化在训练初期的梯度表现更好，能带来更快的训练速度。因此，在很多现代的 Transformer 实现中，Pre-Norm 是更常见的选择。

e.e. 注意力缩放

我们知道，注意力分数的计算核心是 Query 向量 qq 和 Key 向量 kk 之间的点积 qTkq^T k。当向量的维度 dd 变得很大时（比如 Transformer 中常见的512维），两个随机向量的点积结果的量级也可能会变得非常大。点积的方差会与 dd 成正比。

而 Softmax 函数对非常大或非常小的输入值非常敏感。如果输入值很大，Softmax 的输出会趋近于一个独热向量。在这种情况下，梯度会变得极其微小，几乎为0，也就是梯度饱和。一旦发生这种情况，模型就无法通过反向传播进行有效学习了。

研究中发现，点积的量级大致会随着维度的平方根 d\sqrt d 增长。因此，为了抑制这种不受控制的增长，我们在计算 Softmax 之前，将所有的点积结果都除以这个缩放因子 d\sqrt d：

α=softmax(x1:nQK⊤x1:n⊤d)∈Rn×n\alpha = \text{softmax}\left(\frac{x_{1:n}QK^\top x_{1:n}^\top}{\sqrt{d}}\right) \in \mathbb{R}^{n \times n}

f.f. Transformer 编码器

i.i. 基本概念

编码器的主要工作是接收一个完整的序列（比如一整个句子），然后为序列中的每一个词生成一个富含上下文信息的、深度的表示（representation）。

  • 输入： 一个单独的序列 w1:nw_{1:n}。
  • 特点： 编码器不使用未来掩码。在处理任何一个词时，它都可以同时“看到”句子中所有其他的词，包括它前面和后面的词。这使得编码器能够构建双向的上下文表示，信息是完全流通的。

编码器架构非常适合那些不需要自回归式生成文本，而是需要对整个输入序列进行整体理解的任务，例如：

  • 文本分类（判断一篇文章的情感）
  • 命名实体识别（从一句话中识别人名、地名）
  • BERT 这类预训练语言模型
ii.ii. 工作流程

编码器的工作流程如下：

alt text

  1. 输入与嵌入: 将输入的词序列 w1:nw_{1:n} 通过嵌入层 EE 转换为向量序列 x1:nx_{1:n}。
  2. 位置编码: 在向量序列中加入位置信息。
  3. 编码器模块堆叠: 将处理过的向量序列送入一个由 NN 个编码器模块堆叠而成的结构中。每个模块的参数是独立的，前一个模块的输出是后一个模块的输入。

经过 NN 个模块处理后，编码器最终输出一个与输入序列等长的向量序列。此时，每个向量都包含了整个输入序列的上下文信息。

如果需要将输出转换为概率分布（比如在 BERT 的遮盖语言模型任务中），可以在最终的输出后面再接一个线性层和一个 Softmax 层。

iii.iii. 内部结构

每一个编码器模块都由两个核心层组成：

  1. 多头自注意力层: 负责捕捉序列内部各个词之间的关系。
  2. 前馈网络层: 一个简单的全连接神经网络，进行非线性变换。

在这两个子层的后面，都紧跟着一个 Add&Norm 操作，以帮助优化梯度流和稳定训练。

g.g. Transformer 解码器

解码器的主要任务是自回归地生成文本。所谓“自回归”，就是一次生成一个词，并且在生成下一个词的时候，只能参考已经生成好的词。

  • 与编码器的核心区别： 解码器与编码器的结构非常相似，但有一个根本性的不同：解码器的自注意力层中使用了未来掩码。

alt text

h.h. 编码器-解码器结构

Transformer 编码器-解码器 (Encoder-Decoder) 结构是为经典的序列到序列 (Sequence-to-Sequence, Seq2Seq) 任务设计的，比如机器翻译或文本摘要。它将编码器和解码器连接起来协同工作。

当我们希望对输入有双向的、完全的理解（比如阅读整篇文章），但又需要自回归地生成输出（比如逐字生成摘要）时，这个架构就非常有用。它有如下的优缺点：

  • 优点： 在中等模型规模下，其性能通常优于仅解码器模型。
  • 缺点： 需要将模型参数分配给编码器和解码器两部分。目前，许多规模最大的 Transformer 模型都选择了更简洁的仅解码器架构。
i.i. 工作流程

alt text

  1. 编码器部分：
    • 接收输入序列 x1:nx_{1:n}（例如，待翻译的德语句子）。
    • 利用双向上下文理解能力，对输入序列进行编码，生成一系列富含上下文的表示 h(x)h(x)。
  2. 解码器部分：
    • 接收目标序列 y1:my_{1:m}（例如，已经翻译出的英语部分）。
    • 在解码器的每个模块中，除了原有的带掩码的自注意力层之外，还增加了一个新的关键层：交叉注意力。
ii.ii. 交叉注意力

交叉注意力层是连接编码器和解码器的桥梁。它和自注意力略有不同：

  • 在自注意力中，QQ, KK, VV 都来自同一个序列（自己和自己比对）。
  • 而在交叉注意力中，QQ, KK, VV 来自不同的序列：
    • QQ: 来自解码器的序列 h(y)h(y)。
    • KK 和 VV: 来自编码器的最终输出 h(x)h(x)。

可以想象解码器在准备生成下一个英语单词时，会先根据已生成的内容形成一个“问题”或“查询”，然后拿着这个问题去 “查阅”编码器对整个德语输入句子理解后的“知识库”，看看德语句子的哪个部分对生成下一个英语单词最重要，然后把那部分信息重点拿过来参考。