Transformer 详解 | Pytorch 代码解读

前言

  看了许多 Transformer 的文章感觉还是有点抽象，本文将根据 Pytorch 代码从原始输入到最终输出完整走一遍。

代码地址：https://github.com/hyunwoongko/transformer
论文地址：Attention Is All You Need

跑通代码

（1）环境配置

conda create -n transformer python=3.8
conda activate transformer

pip install torch==1.7.0+cu110 torchvision==0.8.1+cu110 torchaudio==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install torchtext==0.8.0
pip install spacy
python -m spacy download de_core_news_sm
python -m spacy download en_core_web_sm

（2）修改部分代码

<1> util/data_loader.py line6
from torchtext.legacy.data import Field, BucketIterator
from torchtext.legacy.datasets.translation import Multi30k
改为
from torchtext.data import Field, BucketIterator
from torchtext.datasets import Multi30k

<2> models/model/transformer.py line55
trg_mask = trg_pad_mask & trg_sub_mask
改为
trg_mask = trg_pad_mask & trg_sub_mask.bool()

（3）配置数据集

数据集自动下载基本失效，可使用以下网址手动下载
https://raw.githubusercontent.com/neychev/small_DL_repo/master/datasets/Multi30k/training.tar.gz
https://raw.githubusercontent.com/neychev/small_DL_repo/master/datasets/Multi30k/validation.tar.gz
https://raw.githubusercontent.com/neychev/small_DL_repo/master/datasets/Multi30k/mmt16_task1_test.tar.gz

解压后将文件放到指定目录，目录结构如下（测试集需修改文件名）

.data
└── multi30k
    ├── test2016.de
    ├── test2016.en
    ├── test2016.fr
    ├── train.de
    ├── train.en
    ├── val.de
    └── val.en

Data & Input

  原始数据就是两个文本文件，每行一句话，对应了原文和翻译，下面介绍如何从文本数据转化成模型的输入矩阵数据。

（1）分词

  spacy 库以及下载的 de_core_news_sm 和 en_core_web_sm 就是用于对两种语言的句子进行分词，从结果上来看大致就是每个单词和标点为一个词。

（2）构建词典

  分词以后对每个词出现的次数进行统计，并按照词频从高到低排序加入词典。具体来说，词典最开始为 4 个特殊符号 '<unk>', '<pad>', '<sos>', '<eos>'，后面就是按词频排序的单词，并且词频小于 2 的词不会加入词典。

（3）构建 Batch

['a', 'blond', 'woman', 'in', 'blue', 'clothes', 'is', 'playing', 'a', 'cello', '.']
['<sos>', 'a', 'blond', 'woman', 'in', 'blue', 'clothes', 'is', 'playing', 'a', 'cello', '.', '<eos>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>']
[2, 4, 122, 14, 6, 29, 294, 10, 37, 4, 1767, 5, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]

  以一个 Batch 中的一个句子为例，第一行是原始分词的结果，第二行对其进行 pad。具体来说就是在句子的开头和结尾添加 <sos> 和 <eos> 特殊符，分别代表 Start Of Sentence 和 End Of Sentence，<pad> 则是对句子填充，使一个 Batch 中每个句子长度相等。第三行就是将每个词用词典中的索引表示，当这个词不在词典中时，也就是用未知的 key 去索引会默认返回 0，对应 <unk>。

Train Inference

Embedding

（1）Token

  模型的输入为 [B,L] 大小的矩阵，其中 B = 128 为 BatchSize，L 为当前 Batch 每个句子的长度，数值为每个词在词典中的索引。
  Embedding 使用 torch.nn.Embedding，其实就是有个 [N,D] 大小的可学习参数，其中 N 为词典的大小，D = 512 为词向量的维度，即每一行都是一个词向量；将输入对这个参数矩阵进行索引就得到了句子中每个词的词向量。

（2）Position

  位置编码用于体现单词在句子中的先后顺序，这里采用了人工生成的方式，直接看公式和代码。
$$\begin{array}{rl}P{E}_{(pos,2i)}& =\sin \left(pos/10000^{2i/d_{\text{model }}}\right)\\ P{E}_{(pos,2i+1)}& =\cos \left(pos/10000^{2i/d_{\text{model }}}\right)\end{array}$$

class PositionalEncoding(nn.Module):
	def __init__(self, d_model, max_len, device):
		super(PositionalEncoding, self).__init__()
		"""
		max_len = 256, 代表句子最大的长度
		d_model = 512
		self.encoding 是固定的, 不参与梯度计算
		pos 对应词的位置
		_2i 对应特征对应的维度
		"""
		self.encoding = torch.zeros(max_len, d_model, device=device)
        self.encoding.requires_grad = False

		pos = torch.arange(0, max_len, device=device)
        pos = pos.float().unsqueeze(dim=1)
        _2i = torch.arange(0, d_model, step=2, device=device).float()
        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))
        
	def forward(self, x):
		batch_size, seq_len = x.size()
		return self.encoding[:seq_len, :]

  分析一下为什么用这种方式做位置编码（个人理解）

• 编码独特性
  • 位置独特性
    下图中第一行绘制了不同位置对应的每个维度的特征值，随着位置数值的增加，图像大致呈右移的效果，这使得不同位置单词的位置特征向量是不同的。
  • 特征维度独特性
    第二行表示了不同维度上，每个位置的特征值，通过使用 $2i/d_{\text{model }}$ 使得每个不同的维度频率是不同的，通常来说不同的特征维度理解为对事物不同维度的特征描述，这种编码方式使得不同维度频率必定不同，对应描述的特征不同。
• 平滑性
  正弦余弦函数在小范围内变化平滑，使位置信息相对连续。
• 扩展性
  对于更长的句子或者更高的特征维度此方式都能 hold 住。
• 周期性
  感觉上正弦余弦函数的周期性一定会有所作用。无论是一个位置的不同特征维度，还是一个特征维度在不同位置上，绝对数值都可能因为周期性产生重复。猜想这样可以使得网络从相对差异来挖掘位置特征，而不是关注绝对的数值。

Encoder & Decoder

  论文原图画的很好，但是细节方面还是有些缺失，后文将描述从输入到输出的每一个计算步骤和细节，与论文原图可以一一对应。

1. Encoder

  下图按以下参数为例进行绘制，当前 Batch 句子长度设为 32，Batchsize 为 128，$d_{\mathrm{model}}=512$，Heads 数量为 $h=8$。

  输入通过 3 个全连接层得到 QKV，QK 计算 Attention Score 作为 V 的权重矩阵，通过矩阵乘法进行加权求和再通过 1 个全连接层就得到 Self-Attention Output，对应论文原图中 Muti-Head Attention 的输出。后续就是残差连接、Add & Norm、以及由 2 个全连接层构成的 Feed Forward。

2. Decoder

  这里假定对应的目标语言输入为 $128\times 27\times 512$，实际上原始输入的句子长度为 28，去掉了句子中的最后一个单词。翻译任务是输出每个单词后面一个单词是什么，而一个 Batch 中最长的句子末尾为 <eos>，很明显不需要预测 <eos> 后面是什么词。

  假设任务是想将英语翻译为德语，将英语记作 Source，德语记作 Target。Self-Attention 的计算方式和前面相同，Decoder 先对 Target 做 Attention，然后将输出的 Target 作为 Q，Encoder 的输出 Source 作为 K 和 V 再做一次 Attention。
  关于多个 Encoder 和 Decoder 模块的堆叠论文原图可能有点歧义，Source 会先通过多个 Encoder 得到一个最终结果，每一个 Decoder 模块都会使用这个结果作为 KV 的输入，而不是使用对应第 N 个 Encoder 模块的输出。

  Decoder 最后的输出接上一个全连接层得到 $128\times 27\times 7853$ 即对词典中所有单词计算分类概率，Loss 计算使用 torch.nn.CrossEntropyLoss。

3. Mask

  虽然在论文原图中只有 Decoder 模块中的第一个 Attention 是带 Mask 的，但实际上在代码中每个 Attention 都会用到 Mask，只不过 Masked Muti-Head Attention 中的 Mask 意义有所不同。在图中我将其标记为 Source Mask 和 Target Mask。
  在说明 Mask 之前需要先梳理一下 Attention。这里忽略 Batchsize 和 Head 维度，$Q[32，64]$ 与 $K^T[64,32]$ 做矩阵乘法得到 $[32,32]$，Softmax 的维度是 -1，意味着每一行是和为 1 的概率值，可以看作句子中每个单词对应该行对应的单词的重要性。
  Attention @ V 的矩阵乘法以结果中的左上角第一个元素进行说明，它由 Attention 的第一行 $A_{0,:}$ 与 V 的第一列 $V_{:,0}$ 点积，也就是每个元素相乘再相加得到。$A_{0,:}$ 代表了句子中每个单词对于第一个单词的重要性，$V_{:,0}$ 为每个单词在第一个维度上的特征，加权求和得到了第一个单词在第一个维度上的特征。

  Mask 的具体作为直接结合代码来看更方便。先看 Mask 的构造，Source Mask 的大小为 $128\times 1\times 1\times 32$，每个句子中 <pad> 部分为 False 其余为 True；Target Mask 的大小为 $128\times 1\times 27\times 27$，除了 <pad> 部分为 False 外，主对角线上方也为 False。

def make_src_mask(self, src):
    src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)
    return src_mask

def make_trg_mask(self, trg):
    trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(3)
    trg_len = trg.shape[1]
    trg_sub_mask = torch.tril(torch.ones(trg_len, trg_len)).type(torch.ByteTensor).to(self.device)
    trg_mask = trg_pad_mask & trg_sub_mask.bool()
    return trg_mask

torch.tril(torch.ones(3, 3))
>>>
[[1, 0, 0],
 [1, 1, 0],
 [1, 1, 1]]

  Mask 的使用如下，将计算得到的 Attention 在 Mask 为 False 的部分变为一个很大的负数，那么在计算 Softmax 时，这个部分就会非常接近 0，达到了遮住这个单词的效果。Source 很好理解，就是遮住了 <pad>；而 Target 是因为在实际翻译的时候只能看到当前输出单词之前的内容。

score = (q @ k_t) / math.sqrt(d_tensor)
if mask is not None:
    score = score.masked_fill(mask == 0, -10000)
score = self.softmax(score)

4. Scale

  举例说明 Attention 计算公式中的 Scale 即 $\sqrt{d_k}$ 的作用，下面直接呈现不使用 Scale 和使用 Scale 得到的某一行 Attention Score 值，可以看出使用后可以大幅缩小概率之间的差异，而网络训练的目的是让 Softmax 的结果差异尽可能增大，或者说突出重要单词的权重，降低其他单词的权重，因此如果不使用 Scale，数值本身的差异已经很大，就会使网络难以训练。
  从数学上看，Softmax 主要使用指数函数 $e^x$，当特征维度较高时，点积运算可能更容易出现较大的 $x$，使用 Scale 对数值进行缩放则可以有效缩小概率值差异。
$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$

[5.6218e-10, 1.0992e-11, 4.7089e-09, 3.3459e-08, 9.1799e-04, 3.3813e-11,
 2.3822e-03, 9.8182e-02, 1.5732e-02, 3.8952e-06, 1.4277e-05, 4.1114e-09,
 7.1132e-01, 2.3731e-02, 2.5220e-02, 1.7592e-11, 2.1984e-09, 3.9202e-10,
 7.4787e-12, 7.0049e-06, 2.5679e-10, 1.2249e-01, 1.1993e-11, 1.5328e-10,
 3.9167e-10, 5.0345e-07]

[0.0100, 0.0061, 0.0130, 0.0166, 0.0597, 0.0070, 
 0.0672, 0.1070, 0.0851, 0.0302, 0.0355, 0.0128, 
 0.1371, 0.0896, 0.0903, 0.0065, 0.0118, 0.0095,
 0.0058, 0.0324, 0.0091, 0.1100, 0.0062, 0.0085, 
 0.0095, 0.0233]

5. LayerNorm & BatchNorm

  两者都是对数据以减去均值再除以标准差的方式做正则化。区别在于对数据的选取，BatchNorm 通常用在 CNN 中，假设输入形状为 $b\times [h\times w]\times d$，对于每个特征维度选取数据（所有样本的所有像素在某个维度上的特征值），即对每个 $b\times [h\times w]\times 1$ 做正则化；LayerNorm 在 NLP 中假设输入为 $b\times l\times d$，对于每个 Batch 选取数据（某个样本所有单词的所有特征值），即对每个 $1\times l\times d$ 做正则化。
  NLP 中使用 LayerNorm 有许多优点。因为输入的序列（句子）长度不一，不像 CV 中输入图像会 Resize 到统一大小，更贴合 NLP 的数据。对于每个样本计算均值方差使得训练和推理阶段 LayerNorm 的行为一致，而 BatchNorm 在训练时计算 Mini-Batch 的均值方差，推理时计算所有样本的均值方差。
  BatchNorm 会受 Batchsize 的影响，过小的 Batchsize 可能会导致不稳定。但在 CNN 中通常会认为每个卷积核抽取不同形式的特征例如颜色、纹理、边缘，对特征维度进行正则化那么大家都是一个特征，只是在绝对数值上做调整，而相对关系保留。如果是 LayerNorm 那么就会让一张图像的不同特征做正则化，而颜色的特征值和边缘的特征值一起算均值方差感觉是不合理的。