0%

Self-Attention 机制与 Transformer 模型

Transformer 模型是由 Google 发表的《Attention is All You Need》论文提出的。它是从 Seq2Seq 模型演进而来的。Transformer 利用 Self-Attention 机制和 Positional Encoding 使得模型在抛弃传统 RNN 结构的同时,既能使用 Attention 的功能也能具备获取句子顺序信息,同时还能并行训练数据。

1. Transfromer 模型结构

Transformer 和 Seq2Seq 一样是由 Encoder 和 Decoder 两部分组成。Encoder 包含 Multi-Head Attention,Add & Norm 和 Feed Forward 层。而 Decoder 同样也包含这三层,只不过是在组合上有些不同而已。Transformer 的输入由 Embedding 和 Positional Encoding 组成。Transformer 的输出经过一层 \(Linear\) 层和 \(Softmax\) 层。

2. Transformer 的输入

在整体上了解了 Transformer 的大致结构后,我们从细节上来看看 Transformer 的结构吧。Transformer 中的输入是由 Embedding 和 Positional Encoding 相加得到的。

2.1 Word Embedding

Word Embedding 有很多中方法可以获取。例如可以使用 Word2Vec、Glove 等算法训练得到,甚至可以在 Transformer 中训练得到。

2.2 Positional Encoding

由于 Transformer 不采用 RNN 结构而是使用全局信息,因此缺乏一种表示输入序列中单词顺序的方法。所以我们在 Transformer 的输入中添加额外的 Positional Encoding。Positional Encoding 向量的维度和 Embedding 的维度要一致的。在 “Attention is all you need” 的论文中,Positional Encoder 是用三角函数去计算的: \[ \begin{aligned} PE_{(pos, 2i)} & = sin(pos / 10000^{2i/d_{model}}) \\ PE_{(pos, 2i+1)} & = cos(pos / 10000^{2i/d_{model}}) \end{aligned} \]

其中 \(PE\) 表示 Positional Encoding; \(pos\) 表示单词在句子中的位置;\(d_{model}\) 表示 PE 的维度(与 Embedding 的维度一致);\(2i\) 表示偶数的维度(\(2i \leq d_{model}\));\(2i+1\) 表示奇数维度(\(2i+1 \leq d_{model}\))。

3. Self-Attention

接下来到 Self-Attention 了。Self-Attention 可以说是 Transformer 模型中最核心的思想,也是不可或缺的一部分。

Self-Attention 的具体计算流程图是这样的:

  • 将输入通过 Positional Encoding 和 Embedding 编码并求和运算得到 Transformer 输入 \(X\)
  • Transformer 的输入 \(X\) 乘以 权重矩阵 \(W_Q, W_K, W_V\) 分别得到 \(Q, K, V\) \[ \begin{aligned} Q & = XW_Q \\ K & = XW_K \\ V & = XW_V \end{aligned} \]
  • 计算 \(Q, K\) 的相关性,即计算 \(f(Q, K)\)\(f\) 是相关性函数。函数 \(f\) 可以有很多种,但我们这里使用点乘。即 \(\theta = Q \cdot K^T\)。其他的相关性函数有:
    • 点乘: \(f(Q, K) = QK^T\)
    • 加权点乘: \(f(Q, K) = Q W K^T\)
    • 加和: \(f(Q, K) = tanh(W_1 Q + W_2 K)\)
  • \(\theta\) 除以 \(\sqrt{d_k}\) 得到 \(\beta\),其中 \(d_k\)\(Q\) 或者 \(K\) 的词向量长度。除以 \(\sqrt{d_k}\) 是为了防止当输入信息的维度 \(d_k\) 比较高时,导致 softmax 函数接近饱和区,进而导致梯度消失
  • 使用 \(Softmax\) 函数对 \(\beta\) 归一化为概率分布 \(\alpha\)
  • 最后用 \(\alpha\) 乘以 \(V\) 得到 \(Attention\)\(Z\)

它的完整计算公式是这样的: \[ \begin{aligned} Attention(Q, K, V) & = softmax(\frac{QK^T}{\sqrt{d_k}})V \end{aligned} \]

4. Multi-Head Attention

Multi-Head Attention 其实就是创建 \(N\) 个(论文中使用 8 个) \(Q, K, V\) 来计算不同的 \(Attention\) 值,并将它们组合起来,再乘以 \(W_O\) 权重矩阵得到 \(Z\)。其中 \(Z\) 的维度和输入 \(X\) 的维度要一致,这样才能在后面进行 \(Add\) 操作。

Multi-Head Attention 机制类似于 CNN 中通过多通道机制进行特征选择。因此能让 Transformer 注意到句子不同的信息点,从而捕捉更多信息。 \[ \begin{aligned} MultiHead(Q, K, V) & = Concat(head_1, \cdots, head_n) W^O \\ where \quad head_i & = Attention(QW_i^Q, KW_i^K, VW_i^V) \end{aligned} \]

5. Add & Layer Norm

Add 操作(residule block,也称残差模块)借鉴了 ResNet 模型的结构,主要作用是使得 Transformer 不因多层叠加而导致梯度消失。

同样的,Layer Normalization 也是为了解决 Transformer 多层堆叠而出现的梯度消失问题。Layer Normalization 把送入激活函数之前的数据进行 Normalization 操作,将其归一化为均值为 0 方差为 1 的数据分布,从而避免因输入数据落在激活函数的饱和区而出现梯度消失问题。Layer Normalization 的计算公式是这样的: \[ \begin{aligned} LN(x_i) & = \alpha \times \frac{x_i - \mu L}{\sqrt{\sigma^2_L + \epsilon}} + \beta \end{aligned} \]

6. Feed Forward

Feed Forward 层(前馈神经网络层)比较简单。它是个两层的全连接神经网络,第一层的激活函数是 \(Relu\),第二层不用激活函数 \[ \begin{aligned} FFN(x) & = Relu(xW_1 + b_1)W_2 + b_2 \\ & = max(0, xW_1 + b_1)W_2 + b_2 \end{aligned} \]

其中 \(X\) 是输入,Feed Forward 的输出矩阵的维度与输入 \(X\) 的一致;\(W_1\)\(b_1\) 是第一层网络的权重与偏置;\(W_2\)\(b_2\) 是第二层网络的权重与偏置。

7. Encoder

上面讲了 Transformer 输入、Self-Attention、Multi-Head Attention、Add & Norm 和 Feed Forward。通过这些组件,我们可以组件 Encoder block 和 Decoder block 了。我们先来说说 Encoder block是如何组成的吧。

从图中看出,Encoder block 是由 \(N\) 个(比如 6 个) Encoder 组成的。其中 Encoder 内部组成如图中右边部分所示。首先 Encoder 的输入 \(X\) 乘以权重矩阵 \(W_i^Q, W_i^K, W_i^V\) 得到 \(Q_i, K_i, V_i\),并进入 Multi-Head Attention 计算 \(Attention\)\(Z\)\(Z\) 的维度与输入 \(X\) 的一致。接着进入 Add & Norm 部分进行残差计算和归一化操作,得到的结果的维度也是与输入 \(X\) 的一致。然后再进入 Feed Forward 层进行计算,计算的结果与输入 \(X\) 的一致。最后再进行 Add & Norm 操作。最终 Encoder 输出的结果与输入 \(X\) 的一致。

8. Decoder

接下来是 Decoder block,如下图所示

Decoder block 也是由 \(N\) 个 Decoder 组成。Decoder 内部和 Encoder 有些不同。Decoder 组了两个 Multi-Head Attention 和 Add & Norm,再到 Feed Forward 和 Add & Norm。这两个 Multi-Head Attention 在细节上有些不同。

  • 第一个的 Multi-Head Attention 的 \(Q_i, K_i, V_i\) 都是从 Decoder 的输入 \(X\) 或者上一层 Decoder 的输出,而且加了 Masked 操作。加入 Masked 操作是因为在翻译的过程中是顺序翻译的,即翻译完第 \(i\) 个单词,才可以翻译第 \(i+1\) 个单词。通过 Masked 操作可以防止第 \(i\) 个单词知道 \(i+1\) 个单词之后的信息。Masked 操作是在计算 Self-Attention,进行 \(Softmax\) 之前使用的,也就是在计算 \(f(Q, K)\) 后,乘以 Masked 矩阵。 说到 Masked,不得不说一下 Transformer 中 Masked 操作有两种,一种是 padding mask,另一种是 sequence mask。sequence mask 就是我们刚才所说的。而 padding mask 则作用在每一个 scaled dot-product attention 中,让过长的句子截断,过短的句子补齐,目的是为了解决输入句子的长度不一的问题。
  • 第二个 Multi-Head Attention 的 \(K_i, V_i\) 是来自 Encoder 的输出 \(Z\) 变换而来的,而 \(Q_i\) 是来自 Decoder 的输入 \(X\) 或者上一个 Decoder 的输出变换而来的。后续的计算方法与之前描述的一致。

Transformer 在训练中也使用了 Teaching Forcing 的训练技巧。也就是说,在训练时,Decoder 的输入使用 ground truth。而在推理时,Decoder 的输入使用的是上一个 Decoder 输出的预测结果。比如在推理时,有一个句子 “我想跑步” 翻译成 “I want to run”。在 Decoder 端

  • 第一次预测:输入是 “起始符 + Positional Encoding”,加入来自 Encoder 的信息是 \(C\) (包含 “我想跑步” 的信息向量),输出是 “I”
  • 第二次预测:输入是 “起始符 + I + Positional Encoding”,加入来自 Encoder 的信息是 \(C\),输出是 “want to”
  • 第三次预测:输入是 “起始符 + I + want to + Positional Encoding”,加入来自 Encoder 的信息是 \(C\),输出是 “run”

由此可看出在训练时,Transformer 是并行的,在推理时是串行的。

参考