Transformer的矩阵维度分析和Mask详解

mask是Transformer中很重要的一个概念，mask操作的目的有两个：

• 让padding(不够长补0)的部分不参与attention操作
• 生成当前词语的概率分布时，让程序不会注意到这个词背后的部分

上面的第一个目的分别对应的是普通的Scaled Dot-Product Attention中的mask操作，而后一个目的对应的是Masked Multi-Head Attention中的Masked

至于这两个mask分别是如何进行的，我们下面来一一讲解。但是首先我们要弄明白multi-attention中的矩阵维度变化，transformer是如何训练和测试的
    
    

Multi-Head attention中矩阵维度的变化

• 假设Multi-Head Attention层的输入为$Q\in {\mathbb{R}}^{(N,T_q,d_{model})}$, $K\in {\mathbb{R}}^{(N,T_k,d_{model})}$, $V\in {\mathbb{R}}^{(N,T_k,d_{model})}$。其中$N$是batch_size，$T_q$时$Q$的maxlen，$T_k$是$K$的maxlen，$d_{model}$是最初单个词embedding的向量长度

• 接下来进行的是$h$次线性变换，$h$实际就是多头注意力的头数，假设第$i$次线性变换后会得到$Q_i^{\ast },K_i^{\ast },V_i^{\ast }$，变换方式如下：
  $$\begin{gathered} Q_i^{\ast }=Q{W}_i^Q{W}_i^Q\in {\mathbb{R}}^{(d_{model},d_k)} \\ {K}_i^{\ast }=Q{W}_i^K{W}_i^K\in {\mathbb{R}}^{(d_{model},d_k)} \\ {V}_i^{\ast }=V{W}_i^V{W}_i^V\in {\mathbb{R}}^{(d_{model},d_v)} \end{gathered}$$
  因此，$Q_i^{\ast }\in {\mathbb{R}}^{(N,T_q,d_k)}$, $K_i^{\ast }\in {\mathbb{R}}^{(N,T_k,d_k)}$, $V_i^{\ast }\in {\mathbb{R}}^{(N,T_k,d_v)}$

• 接下来将$Q_i^{\ast },K_i^{\ast },V_i^{\ast }$进行attention的操作我们分别来看：
  $$\begin{gathered} tem{p}_1=\frac{Q_i^{\ast }(K_i^{\ast }{)}^T}{{\sqrt{d}}_k}tem{p}_1\in {\mathbb{R}}^{(N,T_q,T_k)} \\ hea{d}_i=attention(Q_i^{\ast },K_i^{\ast },V_i^{\ast })=softmax(temp)V \end{gathered}$$
  因此我们可以看出来$hea{d}_i\in {\mathbb{R}}^{(N,T_q,d_v)}$

• 然后我们进行的操作是，将$h$次attention得到的操作连接起来，并乘$W^O$矩阵
  $$\begin{gathered} tem{p}_2=concat((hea{d}_1,hea{d}_2,......,hea{d}_h),-1)tem{p}_2\in {\mathbb{R}}^{(N,T_q,h{d}_v)} \\ Multi-Head(Q,K,V)=tem{p}_2{W}^O{W}^O\in {\mathbb{R}}^{(N,h{d}_v,d_{model})} \end{gathered}$$
  因此最终得到的结果的维度是${\mathbb{R}}^{(N,T_q,d_{model})}$。你应该会发现，和原本输入的$Q$是同样的维度

    
    

Transfromer的训练过程

我们首先做一个假设，训练的batch大小为1,暂时忽略掉Positional Encoding

一个batch包含代翻译的句子集$X$和翻译后的句子集（ground_truth）$Y$，其中:
$$\begin{gathered} X\in {\mathbb{R}}^{(1,T_1,d_{model})} \\ Y\in {\mathbb{R}}^{(1,T_2,d_{model})} \end{gathered}$$
$T_1$是代翻译句子的最大长度，$T_2$是ground_truth的最大长度

联系上图，训练过程的encoder下方的Inputs是$X$,decoder下方的Outputs是$Y$

• 因为Feed Forward层不会改变矩阵维度，所以encoder的输出的向量空间其实就是${\mathbb{R}}^{(1,T_1,d_{model})}$
• 与上面同理，Masked Multi-Head Attention的输出所属的向量空间是${\mathbb{R}}^{(1,T_2,d_{model})}$
• decoder的第二层attention并非self-attention，encoder的输出作为$K$和$V$, Multi-Head Attention的输出作为Q，由我们在上一节的结论，这里的输出属于的向量空间应该是${\mathbb{R}}^{(1,T_2,d_{model})}$
• 上述的linear和softmax操作最终将矩阵变为维度为$(1,T_2,vocab\_size)$，vocab_size是翻译的目标语言词表大小

最终softmax输出的矩阵维度是$(1,T_2,vocab\_size)$，$T_2$表示的是groud_truth的maxlen，因此矩阵第二维的每一个向量都代表着这个句子中一个词的概率分布，比如$（1，2，3）$代表的是第一个句子的第二个词是词表第三个词的概率

由上也不难理解为啥transformer的训练过程可以并行

因此transformer其实是对groud_truth的每一个词进一次预测，所以其损失函数就可以是交叉熵函数

    
    

Transformer的句子生成过程

对于测试样例或者没有ground_truth的样本，Transformer生成句子的模式就和上图差不多(真实的Transformer是由6个encoder和6个decoder组成的)

每一个词的生成都需要输入之前生成的内容，这也是为什么测试和评估过程不是并行的原因

    
    

Mask

最后，通过上面的内容，解释一下为什么要mask以及是如何操作的

在transformer中mask操作其实由三种，我们可以简单的分为：

• 对key中padding进行的mask操作
• 对query中的padding进行的mask操作
• Masked Multi-Head中的mask操作

我们依次来进行介绍，在此之前我们进行一个假设$T_1=4,T_2=4,d_k=4,d_v=4$
    

mask矩阵

我们来看一个维度和单个句子一样的矩阵
$$\left[\begin{array}{cccc}1& 1& 1& 1\\ 1& 1& 1& 1\\ 0& 0& 0& 0\\ 0& 0& 0& 0\end{array}\right]$$
以上就是一个mask矩阵，这个mask矩阵表示的是句子后两个词是padding的内容，所以全都是0。我们在进行attention时不应该将这里考虑进attention，当给mask矩阵为0的对应位置替换一个负很大的值后，相应attention的结果就会趋近为0。

这部分参考了代码transformer的实现,接下来的mask的解释主要是基于该代码
    

对K进行mask

首先我们应该知道，对key进行mask的操作实际上是在attention的softmax之前的，参考代码中的mask矩阵是如下形式的
$$\left[\begin{array}{cccc}1& 1& 0& 0\\ 1& 1& 0& 0\\ 1& 1& 0& 0\\ 1& 1& 0& 0\end{array}\right]$$
因为这个mask矩阵起作用的时间点是$Q{K}^T$计算之后，对$Q{K}^T$进行的mask操作，此时由于$K$经过了转置，所以相应的0的位置也变到了列上面来，于是将$Q{K}^T$的响应位置替换为负的很大值，那么$softmax(Q{K}^T)$对应位置的attention权重就会接近0

padding_num = -2 ** 32 + 1
masks = tf.sign(tf.reduce_sum(tf.abs(keys), axis=-1))  # (N, T_k)，经过这个计算后，若是padding的部分就会变为0
masks = tf.expand_dims(masks, 1) # (N, 1, T_k) 在1的位置上增加一个维度
masks = tf.tile(masks, [1, tf.shape(queries)[1], 1])  # (N, T_q, T_k) tile用于扩展倍数，本句表示二维扩展T_q倍
# ones_like表示的是创建一个同样维度全都是1的矩阵
paddings = tf.ones_like(inputs) * padding_num
# tf.where(input, a, b)表示a中input对应位置维0的不变，其余的替换为b位置的数值
outputs = tf.where(tf.equal(masks, 0), paddings, inputs)  # (N, T_q, T_k)

    

对Q进行mask

对Q进行的mask操作其实是最简单的了，因为$Q$也存在padding的位置，在进行一个Multi-Head Attention计算后，就使得原来是0的位置不是0，所以attention输出的这些位置也应该为空，所以只需要在attention计算之后把相应的位置替换为0即可

masks = tf.sign(tf.reduce_sum(tf.abs(queries), axis=-1))  # (N, T_q)
masks = tf.expand_dims(masks, -1)  # (N, T_q, 1)
masks = tf.tile(masks, [1, 1, tf.shape(keys)[1]])  # (N, T_q, T_k)
outputs = inputs*masks

    

Masked Multi-Head Attention中的Mask

这个稍微复杂一点，我们先看看对应的mask矩阵
$$\left[\begin{array}{cccc}1& 0& 0& 0\\ 1& 1& 0& 0\\ 1& 1& 1& 0\\ 1& 1& 1& 1\end{array}\right]$$
这个是在$Q{K}^T$之后，softmax之前的进行的，我们来慢慢分析

• Masked Multi-Head Attention是Self-Attention，因此$Q=K=V$,我们假设句子是这样的(其中$s_1,s_2,s_3,s_4$都是词向量，并假设$T_1=T_2=4$)

$$Q=K=V=\left[\begin{array}{c}{s}_1\\ s_2\\ s_3\\ s_4\end{array}\right]$$

• 接下来进行计算$Q{K}^T$
  $$Q{K}^T=\left[\begin{array}{cccc}{s}_1{s}_1^T& s_1{s}_2^T& s_1{s}_3^T& s_1{s}_4^T\\ s_2{s}_1^T& s_2{s}_2^T& s_2{s}_3^T& s_2{s}_4^T\\ s_3{s}_1^T& s_3{s}_2^T& s_3{s}_3^T& s_3{s}_4^T\\ s_4{s}_1^T& s_4{s}_2^T& s_4{s}_3^T& s_4{s}_4^T\end{array}\right]$$

• mask矩阵作用于$Q{K}^T$上，可以得到(上面程序中用$2^{32}+1$代替$\infty$)
  $$\left[\begin{array}{cccc}{s}_1{s}_1^T& -\infty & -\infty & -\infty \\ s_2{s}_1^T& s_2{s}_2^T& -\infty & -\infty \\ s_3{s}_1^T& s_3{s}_2^T& s_3{s}_3^T& -\infty \\ s_4{s}_1^T& s_4{s}_2^T& s_4{s}_3^T& s_4{s}_4^T\end{array}\right]$$

• 这里省略掉除以${\sqrt{d}}_k$，进行softmax可以得到
  $$\left[\begin{array}{cccc}{a}_{11}& 0& 0& 0\\ a_{21}& a_{22}& 0& 0\\ a_{31}& a_{32}& a_{33}& 0\\ a_{41}& a_{42}& a_{43}& a_{44}\end{array}\right]$$

• 然后和V相乘
  $$\left[\begin{array}{cccc}{a}_{11}{s}_1& 0& 0& 0\\ a_{21}{s}_1& a_{22}{s}_2& 0& 0\\ a_{31}{s}_1& a_{32}{s}_2& a_{33}{s}_3& 0\\ a_{41}{s}_1& a_{42}{s}_2& a_{43}{s}_3& a_{44}{s}_4\end{array}\right]$$

• 那么使用$[a_{21}{s}_1,a_{22}{s}_2,0,0]$来预测第二个词的概率分布的时候，已经没有了第三个词，第四个词的信息，这就是让训练在生成当前词的时候不会注意到之后的词的原因

diag_vals = tf.ones_like(inputs[0, :, :])  # (T_q, T_k)
# 这一句的意思是生成一个上三角矩阵，上三角矩阵用来对decoder的结果进行mask
tril = tf.linalg.LinearOperatorLowerTriangular(diag_vals).to_dense()  # (T_q, T_k)
masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(inputs)[0], 1, 1])  # (N, T_q, T_k)
paddings = tf.ones_like(masks) * padding_num
# 将mask为0处全部变为负无穷
outputs = tf.where(tf.equal(masks, 0), paddings, inputs)

[参考链接]

参考文档：attention is all you need

参考程序:Kyubyong/transformer

如有错误，敬请指正，欢迎交流