理清Transformer输入输出

前言

之前写完一篇笔记以为自己搞得很懂了，其实还是有些小细节没注意到，特别是跟NLP比较有关的内容，就再记录一下Transformer的整个输入输出过程。

padding

不用担心输入的句子不一样长，不用担心预测的时候decoder的输入不是完整的句子，因为在输入decoder和encoder时都会进行补0，把句子长度补到一样，因此序列长度是固定的N。

训练阶段

假设batch_size记为B，序列长度记为N(已进行padding)，每个单词对应的 one-hot 编码维度也即 vocab_size 记为$d_{vs}$。

Encoder

注意：下面的乘法都是一个batch一个batch地乘（就是torch.bmm）。

输入的是B个长度为N的序列，记作$X\in R^{B\times N\times d_{vs}}$；

进行word embedding后把每个单词对应的向量的维度变为 $d_{model}$，此时维度为$R^{B\times N\times d_{model}}$；
加上维度也为$R^{B\times N\times d_{model}}$的postional encoding；

输入multi-head attention：
对于第i个head，分别乘以矩阵$W_i^Q,W_i^K,W_i^V$得到$Q_i,K_i,V_i$，维度分别为$R^{B\times N\times d_q},R^{B\times N\times d_k},R^{B\times N\times d_v}$；
$Q_i,K_i$ 做scaled dot-product attention（$A_i=Q_i{K}_i^T,A_i^{\prime }=softmax(\frac{A_i}{\sqrt{d_k}})$）得到$A_i^{\prime }\in R^{B\times N\times N}$ ；
$A_i^{\prime }$与$Q_i$做乘法（就是加权求和），得到输出 $hea{d}_i$ 的维度为$R^{B\times N\times d_v}$；
合并所有 $hea{d}_i$，先将它们按列拼接在一起，得到 $(hea{d}_1,...,hea{d}_h)\in R^{B\times N\times h{d}_v}$，再乘以矩阵 $W^O\in R^{B\times h{d}_v\times d_{model}}$，得到输出维度为 $R^{B\times N\times d_{model}}$；

论文中，取$d_v=d_k=d_q=d_{model}/h$；

add&norm：
没什么好说的，加上multi-head attention的输入后再layer norm，维度仍为 $R^{B\times N\times d_{model}}$；

feed forward：
这个子层是由2个全连接层和中间一个激活函数组成的。第一个全连接层把维度变为$R^{B\times N\times 2d_{model}}$，第二个全连接层再变回$R^{B\times N\times d_{model}}$；

add&norm：
同前一个；

Decoder

在训练阶段decoder的输入是完整的序列(ground truth)，word embedding部分和positional encoding部分与encoder一致；

进入masked multi-head attention：
decoder在预测下一个输出的时候是不可以看到真实答案的后面的内容的，所以需要把后面的内容遮住。举个例子：翻译答案为"机器学习"，第一个输入[BOS] (begin of sequence的意思)，输出预测’鸡’（注意这里预测错了）；第二次要输入[BOS]和’机’（注意：在训练时前面预测错了，不用管他，直接输正确答案就行），输出预测’器’…
masked的操作方法是在softmax前把attention scores矩阵$A$对应位置的值设为$-\infty$，这样softmax后对应的 attention scores 会变成0，加权求和的时候就不会加上它们。
其它都跟encoder的一样，输出维度为$R^{B\times N\times d_{model}}$；

add&norm：
同前一个；

进入multi-head attention：
这个是连接encoder和decoder的地方，计算方式是相同的，只是这里的 $Q_i$ 是用decoder前面的输出乘以矩阵 $W_i^Q$ 得到的，$K_i,V_i$ 是用encoder一个block的输出乘以 $W_i^K,W_i^V$ 得到的。
输出维度为$R^{B\times N\times d_{model}}$；

add&norm：
同前一个；

feed forward：
同encoder；

add&norm：
同前一个；

Linear & softmax：
这个地方就是要输出每个位置对应词汇表中每个单词的概率。
Linear层要用前面得到的维度为$R^{B\times N\times d_{model}}$的输出，乘以维度为 $R^{B\times d_{model}\times d_{vs}}$的矩阵，得到输出维度为$R^{B\times N\times d_{vs}}$，再进行softmax得到最终的概率矩阵。
最后列中取softmax值最大对应的那个单词就是预测值。

测试阶段

测试阶段跟训练阶段唯一的不同就是在decoder的输入。测试阶段是没有正确答案的，输入的序列就是原原本本的模型的预测值（比如前面举的例子中，预测错的’鸡’是要作为输入的），也因此不需要遮住后面的序列（因为根本没有）。