Transformer的架构理解

前言

本文是结合了B站李沐论文精读系列课程的内容和Transformer论文的内容，再加上自己的整理。
本文主要涉及到Transformer网络结构的LayerNorm、Attention、Embedding、Positional Encoding部分。
感谢李沐大神的视频课帮我对Transformer架构有了更深入的了解。

LayerNorm & BatchNorm

BatchNorm

BatchNorm是对每一个特征做Normalization。

以最简单的二维输入矩阵为例，矩阵的行为batch，列为特征feature。

BatchNorm就是要把每一列（特征）的均值变为0，方差变为1。也就是每一列都减掉均值，再除以它们的标准差。

BatchNorm还会学习到 $\lambda ,\beta$ 两个参数，也就是可以把特征向量通过学习，放成任意方差和均值的向量。

### LayerNorm

LayerNorm是对每一个样本做Normalization，也就是把每一行（样本）的均值变为0，方差变为1。二维情况下如上图右所示。

如果是三维情况下呢？

BatchNorm是对同一个特征做Normalization，这就导致均值和方差是跨样本的，也就是考虑了许多样本，一起求一个均值和方差。在样本长度变化比较大的时候，每次做小批量的均值和方差它的抖动是比较大的。另外，在预测的时候，我们是要把全局的均值和方差记录下来，如果新的预测样本这个序列很长，在训练的时候没有见到过这么长的序列，那么用之前的均值和方差就不合适了。

LayerNorm是对一个样本来做Normalization，这样计算均值和方差的时候都是在这同一个样本内计算，不需要保存全局的均值和方差，不会被其他样本干扰，更为稳定。

Decoder

解码器使用自回归，也就是当前的输入是上一个时刻的输出。也就是说解码器在预测的时候不应该看到之后时刻的输出，而注意力机制每次都能看到整个完整的输入，所以为了避免这种情况，使用Masked Multi-Head Attention，将t时刻后面的输入mask掉。

Attention

注意力函数可以描述为将一个query和一系列key-value对映射到一个输出，其中这些query、key、value和output都是向量。输出output是value的加权和，其中这些权重由query和这些keys的相似度计算而来。也就是说，如果query和某个key的相似度高，那么这个key对应的value的权重就高。

所以，同样的key-value对，由于query的不同，导致values的权重不同，计算出来的output也不同，这就是注意力机制。

不同的相似度函数会产生不一样的注意力版本，在Transformer中使用的是最简单的注意力机制，Scaled Dot-Product Attention。

Scaled Dot-Product Attention

query和key的长度都是 $d_k$ (query和key等长，当然也可以不等长，不等长是由别的方法可以计算)，value的长度为 $d_v$ 。计算query和所有keys的内积，再除以 $\sqrt{d_k}$ ，然后应用一个softmax函数来得到values的权重。

在实现中，通过将所有的query打包成一个矩阵Q来实现并行计算一系列query。keys和values也被打包为矩阵K和V。公式如下：
$$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V\text{\hspace{1em}(1)}$$
解释：

• 一个query对所有的key计算，得到n个值，然后将这n个值经过一个softmax函数得到n个非负的和为1的权重值。把这个权重作用在values上面得到输出。

• 假设有n个query，m个key-value对，那么矩阵Q的形状为$(n,d_k)$ ，矩阵K的形状为 $(m,d_k)$ ，所以 $Q{K}^T$ 的形状为 $(n,m)$ 。矩阵V的形状为 $(m,d_v)$ ，所以最终的输出矩阵的形状为 $(n,d_v)$ 。其中的每一行代表一个query对应的输出。使用两次矩阵乘法，就得到了所有的注意力，方便并行计算。

• 时间复杂度：

  • 矩阵乘法运算 $Q{K}^T$ ：$O(n\times d_k\times m)=O(n^{2}d)$
  • 对$(n,m)$的矩阵进行softmax运算：$O(mn)=O(n^2)$
  • 与V做矩阵乘法运算：$O(n\times m\times d_v)=O(n^{2}d)$

  所以最终self-attention的时间复杂度为$O(n^{2}d)$。参考知乎文章

有两种不同的注意力机制被广泛使用，一种是加型注意力机制，另一种是内积型注意力机制。内积型就是文章中用到的，只不过除以了一个$\sqrt{d_k}$ ，因此其名为Scaled。而加型注意力机制使用一个只含有一个隐藏层的feed-forward网络，可以处理query和key不等长的情况。这两种注意力机制差不多，但是内积型的计算更快更高效，因此采用这种。

那么为什么要额外除以一个$\sqrt{d_k}$ 呢？

当$d_k$ 的值比较小的时候，除不除这个 $\sqrt{d_k}$ 影响并不大，但是当 $d_k$ 的值比较大的时候，也就是query和key这两个向量长度比较长的时候，内积计算出来的某些值的绝对值就会比较大，这时经过softmax之后，这些值就会更加接近1，其它的值就会接近0，也就是这些值向两极分化。由于softmax的目标就是让置信的更加接近1，不置信的更加接近0，现在已经符合了这种情况，就会导致在计算梯度的时候梯度很小，导致训练不动。除以一个 $\sqrt{d_k}$ 就会让这个梯度不会变得太小，有利于训练。

Mask如何实现？

mask主要是用来避免第t时刻看到以后时刻的东西。

假设query和key等长，并且时间上能对应起来。那么对于第t时刻的query $q_t$ ，应该只使用 { k 1 , . . . , k t − 1 } \{k_1,...,k_{t-1} \} {k1​,...,kt−1​} 来做计算，但是在注意力计算的过程中，实际上是用到了 { k 1 , . . . , k t − 1 , k t , . . . , k d k } \{k_1,...,k_{t-1},k_t,...,k_{d_k} \} {k1​,...,kt−1​,kt​,...,kdk​​} 所有的key。为了消除 $k_t$ 及后面的key的影响，在计算softmax之前，将t时刻后面所有的计算结果更换为一个很大的负数，比如 $-1e10$ ，这个大负数经过了softmax之后就会变成0，这样对应的t时刻及之后的value对应的权重就为0，这样就只考虑了t时刻之前的key。

Multi- Head Attention

把query、key、value通过可学习的线性投影到 $d_k、d_k、d_v$ 维，投影出来的结果再计算Attention，得到一个 $d_v$ 维的输出。这样的过程进行$h$ 次，将得到的 $h$ 个输出拼接在一起，再进行一个线性投影，得到一个 $d_v$ 维度的最终结果。

多头注意使得模型能够联合注意不同位置上不同表示子空间的信息。
$$\begin{gathered} MultiHead(Q,K,V)=Concat(hea{d}_1,...,hea{d}_h)W^O \\ wherehea{d}_i=Attention(Q{W}_i^Q,K{W}_i^K,V{W}_i^V) \end{gathered}$$
其中的投影是参数化的矩阵，$W_i^Q\in {\mathbf{R}}^{d_{model}\times d_k},W_i^K\in {\mathbf{R}}^{d_{model}\times d_k},W_i^V\in {\mathbf{R}}^{d_{model}\times d_v},W^O\in {\mathbf{R}}^{h{d}_v\times d_{model}}$

Point-wise Feed-forward Networks

就是图中的Feed Forward层。其实是一个MLP，只不过是作用在每个向量上，所以称为Point-wise，这个point就是指每个向量，以翻译任务为例，每个point就是句子中的每个单词向量。对所有的单词向量使用同一个MLP，这个MLP含有单隐藏层，它将输入由512维扩大到2048维，然后输出层再将2048维缩小到512维。
$$FFN(x)=max(0,x{W}_1+b_1)W_2+b_2\text{\hspace{1em}(2)}$$

Embedding and Softmax

Embedding: 就是对输入的每一个单词(token)，学习一个长度为d的向量来表示它。

decoder的输出结果再进入一个线性层和一个Softmax层，转换为下一个单词(token)的预测概率。在Transformer模型中，encoder和decoder的两个Embedding和decoder最后的线性层共用权重矩阵，方便训练。在Embedding层，将这些权重乘以 $\sqrt{d_{model}}$ ，这里 $d_{model}=512$ ，目的是为了使输出值的尺度和接下来的Positional Encoding输出值差不多。

Positional Encoding

Attention本身没有时序信息，它在输入里面加入位置信息。例如在单词的Embedding中使用512维的向量表示，那么这里用一个512维向量表示这个单词的位置(0,1,2,3…)，将这个位置向量与Embedding向量相加。