【Transformer 1】Attention Is All You Need 学习笔记

论文的强大之处在于，提出了一种全新的模型范式 Transformer，用于对序列数据建模。Transformer 的核心科技是 Attention Mechanism，包括 Self-Attention（SA）、Multi-Head Attention（MHA）。虽然注意力机制这个概念不是该论文的首创，但将其发扬光大确实是它的功劳。

网上大佬们解读 Transformer 和注意力机制的文章很多，讲解的很详细很精彩，本人也是看了很多文章才对这些概念有了一些理解。本章结只记录本人学习后的一些理解，原理和计算方面不会说的很细（因为也理解不了很细），详细原理还请参考其他大佬们的文章。

1. 模型结构

Transformer 是首个完全基于注意力机制搭建的模型，它解决了：(1) RNN 模型依赖序列数据的时序关系，无法并行化计算；(2) 卷积网络对隐状态的计算量随序列长度增加，难以学习远距离依赖的问题，因此 Transformer 也成为了 GPT 等大语言模型的基础架构，用来处理更长的文本序列数据，也让成百上千亿甚至更大参数规模的模型训练成为可能（虽然训练依然很贵）。

Transformer 的几大部分：

Input：Word Embedding + Positional Encoding
Process： $N\times$ Encoder + $N\times$ Decoder
Encoder：Multi-Head Attention (MHA) + Feed Forword Network (FFN)
Decoder：Masked MHA + MHA + FFN
Output：Linear Transformation + Softmax

2. Positional Encoding

Transformer 使用位置编码来表示 token 在序列中的绝对位置（表示 token 自身的位置）和相对位置（表示 token 与其他 token 的位置关系）信息。编码函数如下：

$PE_{(pos,2i)} = sin(\frac{pos}{10000^{2i/d_{model}}})$

$PE_{(pos,2i+1)} = cos(\frac{pos}{10000^{2i/d_{model}}})$

其中 $pos$ 表示 token 的位置， $i$ 表示 token 中的第 $i$ 个维度， $d_{model}$ 是 Transformer 中每一层的输出维度，文中统一设置为 512。为什么使用这种编码方式，作者给出的解释是他们假设（猜）这种方式有利于模型学习相对位置关系，正弦的形式允许模型将序列外推到比训练序列更长的长度。（作者也实验了可学习的位置嵌入方式，从结果上看没啥差别）对于使用该编码方式更详细的分析可以看 [1] 中的回答。

位置编码向量在输入层与 embedding 相加后送入处理部分。至于为何是相加，可以看大佬 [2] 给出的一种分析。

3. Encoder & Decoder

输入层完成后进入到 Transformer 的主体部分，该部分包括一个 Encoder 和一个 Decoder，分别通过堆叠 $N$ （文中设置为 6）个相同的子块实现。

3.1 Encoder

Encoder 中每个子块如图所示，包括 1 个 MHA 子层和 1 个 FFN 子层，每个子层又跟着一个残差连接和一个 Layer Normalization (LN)。

$y_{en}=LayerNorm(x_{en}+MHA(x_{en}))$

$z_{en}=LayerNorm(y_{en}+FFN(y_{en}))$

3.2 Decoder

Decoder 中每个子块如图所示，包括 1 个 Masked MHA 子层，1 个 MHA 子层和 1 个 FFN 子层，每个子层又跟着一个残差连接和一个 LN。

$w_{de}=LayerNorm(x_{de} + MMHA(x_{de}))$

$y_{de}=LayerNorm(w_{de} + MHA(z_{en}, w_{de}))$

$z_{de}=LayerNorm(y_{de} + FFN(y_{de}))$

Encoder 中的 MHA 的输入均来自 Input Embedding(+PE)，而 Decoder 中的 MHA 的输入来自Encoder 的输出和 Output Embedding(+PE)。Decoder 中使用了带 mask 的 MHA，确保 Transformer 在预测第 $i$ 个位置的单词时只能依赖 $i$ 之前的历史信息，对 $i$ 之后的信息做掩膜处理，防止模型看见未来（防止学习过程中作弊）。

4. 注意力机制

4.1 Scaled Dot-Product Attention

作者将 Transformer 中用到的注意力计算方法叫做 Scaled Dot-Product Attention，具体运算如下：

$Attention(Q,K,V ) = softmax(\frac{QK^T}{\sqrt{d_k}})V$

其中 $Q,K,V$ 是向量矩阵， $Q$ 和 $K$ 中每个向量的维度为 $d_k$ ， $V$ 中每个向量维度为 $d_v$ 。论文中 $d_k=d_v=d_{model}=512$ ，如果是多头注意力， $d_k=d_v=\frac{d_{model}}{N_{head}}=\frac{512}{8}=64$ 。当 $Q,K,V$ 都来自同一组输入，就称为 Self-Attention。

为什么要乘以尺度缩放因子 $\frac{1}{\sqrt{d_k}}$ ？

注意力权重（0-1之间）的计算用到了 softmax，当 $d_k$ 很大时 $Q$ 和 $K$ 做点积的结果也会很大，会将 softmax 推向梯度变化极小的区域，造成梯度消失，不利于模型收敛，乘以 $\frac{1}{\sqrt{d_k}}$ 将 softmax 拉回到梯度变化明显的区域。

假设 $Q$ ， $K$ 中的每个向量 $q_i$ ， $k_i$ 都是独立同分布（均值为0，方差为1）的随机变量，则

$E(QK^T)=E(Q)E(K^T)=0$

$D(QK^T)=D(\sum_{i=1}^{d_k}q_i\cdot k_i^T)=\sum_{i=1}^{d_k}D(q_i\cdot k_i^T)$

$=\sum_{i=1}^{d_k}D(q_i)D(k^T_i)=\sum_{i=1}^{d_k}1\times 1=d_k$

$Q$ 和 $K$ 做点积之后的方差为 $d_k$ ，不管 $d_k=512$ 还是 $d_k=64$ ，这个值都是很大的，当我们乘以一个 $\frac{1}{\sqrt{d_k}}$ ，根据 $D(CX)=C^2D(X)$ 可得

$D(\frac{QK^T}{\sqrt{d_k}})=(\frac{1}{\sqrt{d_k}})^2D(QK^T)=\frac{d_k}{d_k}=1$

点积之后的方差又回到了 1，相当于对数值进行了归一化，保持了数值的稳定性 [7-9]。

4.2 Multi-Head Attention

多头注意力允许模型对一个输入序列同时从不同的特征表示子空间或者说不同视角关注信息 [11]，至于为什么会这样可以看 [10]，从实验结果上看，使用多头确实比单头能产生更好的效果，但也不是越多越好，属于经验超参。

在计算上，多头注意力属于多个单头注意力计算结果的组合，相当于并行计算多个单头

$MultiHead(Q,K,V ) = Concat(head_1,...,head_h)W^O$

$head_i = Attention(QW^Q_i ,KW^K_i ,VW^V_i )$

其中 $W^Q_i \in \mathbb{R}^{d_{model} \times d_k}$ ， $W^K_i \in \mathbb{R}^{d_{model} \times d_k}$ ， $W^V_i \in \mathbb{R}^{d_{model} \times d_v}$ ， $W^O \in \mathbb{R}^{hd_v \times d_{model}}$ ，文中的 baseline 使用了 8 个头，所以 $d_k = d_v =d_{model}/h = 64$ 。由于 $d_k$ ， $d_v$ 减小，总的计算量并没有增加很多。

4.3 Why Self-Attention

作者从三个方面说明了为什么使用 SA，与其它网络类型相比的优势：

时间复杂度：越低越好
最小顺序操作数：用来衡量并行计算量，越小越好
最大路径长度：学习长距离依赖的关键因素，越短越好

先上结果：

（1）时间复杂度

具体的计算过程可以看 [12-14]。对于普通的 SA，时间复杂度是 $O(n^2d)$ ，在 $n<d$ 时 SA 是优于 RNN 的，但 $n\geqslant d$ 时就没有优势了，所以当 $n$ 非常大时考虑了一种受限的 SA，对于序列中任意 token 只考虑与其周围 $r$ 邻域内的 token 做 SA，时间复杂度会降到 $O(rnd)$ 。

对于卷积结构， $O(knd^2)$ 中的 $O(kd)$ 是 $k\times d$ 的卷积核做一次运算的时间复杂度，对于一个 $n\times d$ 的输入，卷积核需要进行 $n$ 次卷积（ $stride=1$ ， $padding=\frac{k-1}{2}$ ），生成1个 $n\times 1$ 的特征向量，时间复杂度为 $O(knd)$ ，共需 $d$ 个卷积核作上述操作，最终输出 $n\times d$ 的特征向量，所以时间复杂度共 $O(knd\cdot d)=O(knd^2)$ 。根据 [12] 和 [14] 中的分析画了一个大致的过程图，输入输出都是 $n\times d$ 。

对于 RNN 的时间复杂度的计算可以看 [14]，过程还是很清楚的。

（2）顺序操作数

序列操作指模型处理序列数据时，必须按输入顺序依次执行的、无法并行化的操作步骤，直接决定了模型的计算效率，尤其是长序列场景。只有 RNN 需要按时间步顺序执行，操作数跟序列长度 $n$ 有关，SA 和卷积都可以并行处理，所以操作数是常数级的。

（3）最大路径长度

最大路径长度表示序列中任意两个 token 之间传递信息所需的最大距离，直接反映模型捕捉长距离依赖关系的能力，路径越短，长距离信息传递越高效。

对于 RNN 比较好理解，处理长度为 $n$ 的序列，第 1 个位置的信息要传递到第 $n$ 个位置，需要 $n-1$ 步。

CNN 采用卷积核滑动的方式扩展感受野。使用大小为 $k$ 的卷积核，1层卷积能覆盖 $k$ 个连续位置，2 层卷积能覆盖 $k^2$ 个位置， $m$ 层卷积能覆盖 $k^m$ 个位置。要覆盖长度为 $n$ 的序列，需要 $l=log_kn$ 层卷积。因此，位置 $i$ 与位置 $j$ 之间的信息传递需经过 $log_k(\left | j-i \right |)$ 步。

SA 是直接计算任意两个位置之间的关联，最大路径长度为 1。

受限的 SA 中 token 的信息传递受到 $r$ 的制约，假设一个 $n=12$ 的序列 $\left \{ t_1,t_2,...,t_{12}\right \}$ ， $t_1$ 与 $t_{12}$ 传递信息，必须要经过中间节点 $t_4,t_7,t_{10}$ ，需要 $n/r$ 步。

综合来看，SA 不仅能加速计算，而且能捕捉长距离依赖，计算量也有一定优势（可以通过分词方法实现 $n<d$ ），补齐了 RNN 和 CNN 的不足。

4.4 Applications of Attention in Transformer

Transformer 中三种形式的 Attention：

（1）Cross Attention between Encoder and Decoder

Decoder 接收 Encoder 输出的 Key 和 Value，与自身生成的 Query 做注意力计算，使得Decoder 在生成下一个单词预测时能充分利用到源序列数据的上下文信息。

（2）Self-Attention in Encoder

qkv 来自同一个序列，序列中每一个 token 可以获取到其它所有 token 的信息。

（3）Masked Self-Attention in Decoder

关于使用 Masked SA 的原因可以看 [16] 的分析，训练阶段比较好理解，模型在预测第 $i$ 个位置的单词时只能综合 $i$ 之前的历史信息做预测，而使用 SA 的 token 之间是全连接的关系，第 $i$ 个位置的 token 包含了 $i$ 之后的所有 token 信息，如果不使用 mask，模型根据未来“预测”现在，这显然不合理。至于预测阶段为什么使用 mask，还是看大佬们的分析吧，本人也是似懂非懂，不敢瞎说。