Attention is all you need记录

Scaled Dot-Product Attention

$Q\in {\mathbb{R}}^{l_q\times d_q},K=V\in {\mathbb{R}}^{l_k\times d_k}$，一般来说K，V是一样的，但也可以不同。在机器翻译(Encoder-Decoder Attention)中，Q为目的语言(Target)， K，V是源语言(Source)。

1. MatMul & Scale: $W_{atten}=\frac{Q\cdot K^T}{\sqrt{d_k}}\in {\mathbb{R}}^{l_q\times l_k}$
2. Mask: Padding Mask and Sequence Mask(optional)
   $$W_{atten}=Mask(W_{atten})$$
3. Softmax: $W_{atten}=softmax(W_{atten})$
4. MatMul: $output=W_{atten}\cdot V\in {\mathbb{R}}^{l_k\cdot d_k}$

本质上，将$V$看成$l_k$个维度为$d_k$的列向量，每一个列向量都代表一个词嵌入(token embedding)，attention其实就是将句子中的每一个词进行线性加权。而权重就是根据$Q$对$K$的计算(具体为Scaled Dot-Production)得来的。

Padding Mask

这个很好理解，在自然语言处理中，通常情况下句子都是不是等长的，但常用的深度学习框架(TF, Pytorch等)都只能处理等长的句子(实际上可以处理变长句子，例如pytorch的PackedSequence，但这不是一般做法)。为了解决这个问题，一般会使用0来进行填充(padding)，这些pad token没有实际意义。在使用注意力机制(Attention)的时候，对于这些token不应该分配注意力，因此需要将其掩盖(Mask)。
具体做法也很简单，在进行softmax之前(上节第2步)，对于padding token所在位置，加上一个很大的负数，这样进行softmax之后，对应位置就变为0了。

  def get_padding_masks(self, query, key):
    '''
    :param Query: # B*n, L_q, d_k/n， where n is num_heads
    :param Key: # B*n, L_k, d_k/n
    :return:
    '''
    neg_large_num = -2 ** 32 + 1
    l_q = query.size(1)
    masks = t.sign(t.sum(t.abs(key), -1)).unsqueeze(1)  # B*n * 1 * L_k
    masks = masks.repeat(1, l_q, 1)  # B*n * L_q * L_k
    paddings = t.zeros_like(masks).fill_(neg_large_num)
    return t.where(masks == t.zeros_like(masks), paddings, t.zeros_like(paddings))

Sequence Mask

所谓Sequence Mask，就是掩盖未来的信息，而止保留当前位置之前的信息。只有Decoder的self-attention会用到Sequence Mask
我们知道Transformer有比LSTM更强的上下文整合能力，那么这个能力体现在哪呢？答案就是self-attention，在一个句子中任意位置的token都可以互相“交互”（时序问题在输入端有Positional Encoding，这里不详细解释）。注意：只有在Decoder端的self-attention中需要使用(对应论文结构图中的Masked Multi-head Attention)。而在Decoder端self-attention和Encoder-Decoder Attention(该部分本质上是Decoder端target token对Encoder端进行Query，不涉及target token与未来时序的token进行交互)中不需要使用。
Eecoder端（source）的所有信息都是透明的，在Decoder端，由于目的就是预测token，不能将未来时序的token都“暴露”出去，因此需要进行“Mask”。
具体做法就是在第1步算出$W_{atten}$之后，将矩阵上半部分赋值为0，形成一个下三角矩阵。

  def get_sequence_masks(self, attn_weight):
    '''
    :param attn_weight: (B*n, L_q, L_q)
    :return:
    '''
    neg_large_num = -2 ** 32 + 1
    masks = t.zeros_like(attn_weight).fill_(neg_large_num)
    return t.tril(masks, diagonal=0)

完整代码

import torch as t
from torch import nn


class Transformer(nn.Module):
  def __init__(self, embed_dim, vocab_size, hidden_dim=64, num_heads=8, num_layers=6):
    super(Transformer, self).__init__()
    self.num_layers = num_layers

    self.pe = PositionalEncoding(embed_dim)
    self.input_linear = nn.Linear(embed_dim, hidden_dim)
    self.multi_head_attention = MultiHeadAttention(embed_dim, hidden_dim, num_heads)
    self.layer_norm = LayerNormalization(hidden_dim)
    self.feed_forward = nn.ModuleList([nn.Linear(hidden_dim, hidden_dim)] * num_layers)
    self.activation_function = t.tanh
    self.output_linear = nn.Linear(hidden_dim, vocab_size)

  def forward(self, query, key=None, value=None, mode="e"):
    query = self.pe(query)
    query = self.input_linear(query)
    if mode.lower() in ["e", "encoder", "encode"]:
      for i in range(self.num_layers):
        outputs = self.layer_norm(query + self.multi_head_attention(query, query, query))
        outputs = self.layer_norm(outputs + self.activation_function(self.feed_forward[i](outputs)))
        query = outputs
      return query
    elif mode.lower() in ["d", "decoder", "decode"]:
      for i in range(self.num_layers):
        query = self.layer_norm(query + self.multi_head_attention(query, query, query, True))
        outputs = self.layer_norm(query + self.multi_head_attention(query, key, value))
        outputs = self.layer_norm(outputs + self.feed_forward[i](outputs))
        query = outputs
      return t.softmax(self.output_linear(query), -1)


class MultiHeadAttention(nn.Module):
  def __init__(self, embed_dim, hidden_dim, num_heads):
    super().__init__()
    assert hidden_dim % num_heads == 0
    self.embed_dim = embed_dim
    self.hidden_dim = hidden_dim
    self.num_heads = num_heads
    self.linear_q = nn.Linear(self.hidden_dim, self.hidden_dim)
    self.linear_k = nn.Linear(self.hidden_dim, self.hidden_dim)
    self.linear_v = nn.Linear(self.hidden_dim, self.hidden_dim)
    self.linear_out = nn.Linear(self.hidden_dim, self.hidden_dim)
    self.activate_function = t.tanh

  def forward(self, query, key, value, sequence_padding=False):
    batch_size = key.size(0)
    Q = self.activate_function(self.linear_q(query))
    K = self.activate_function(self.linear_k(key))
    V = self.activate_function(self.linear_v(value))

    split_size = self.hidden_dim // self.num_heads
    Q_ = t.cat(t.split(Q, split_size, dim=-1), dim=0)  # B*n, L, d_k/n
    K_ = t.cat(t.split(K, split_size, dim=-1), dim=0)
    V_ = t.cat(t.split(V, split_size, dim=-1), dim=0)

    attention = self.scaled_dot_production(Q_, K_, V_, sequence_padding)
    attention = t.cat(t.split(attention, batch_size, 0), -1)
    return self.linear_out(attention)

  def scaled_dot_production(self, Q, K, V, sequence_padding):
    attn_weight = t.matmul(Q, K.permute(0, 2, 1)) / t.sqrt(t.Tensor([self.embed_dim]))  # B*n * L_q * L_k
    masks = self.get_padding_masks(Q, K)  # B*n * L_q * L_k

    if sequence_padding:
      masks += self.get_sequence_masks(attn_weight)
    attn_weight = t.softmax(attn_weight + masks, -1)
    return t.matmul(attn_weight, V)

  def get_padding_masks(self, query, key):
    '''
    :param Query: # B*n, L_q, d_k/n
    :param Key: # B*n, L_k, d_k/n
    :return:
    '''
    neg_large_num = -2 ** 32 + 1
    l_q = query.size(1)
    masks = t.sign(t.sum(t.abs(key), -1)).unsqueeze(1)  # B*n * 1 * L_k
    masks = masks.repeat(1, l_q, 1)  # B*n * L_q * L_k
    paddings = t.zeros_like(masks).fill_(neg_large_num)
    return t.where(masks == t.zeros_like(masks), paddings, t.zeros_like(paddings))

  def get_sequence_masks(self, attn_weight):
    '''
    :param attn_weight: (B*n, L_q, L_q)
    :return:
    '''
    neg_large_num = -2 ** 32 + 1
    masks = t.zeros_like(attn_weight).fill_(neg_large_num)
    return t.tril(masks, diagonal=0)


class LayerNormalization(nn.Module):
  def __init__(self, hidden_dim):
    super(LayerNormalization, self).__init__()
    self.hidden_dim = hidden_dim
    self.alpha = nn.Parameter(t.ones(1, hidden_dim))
    self.beta = nn.Parameter(t.zeros(1, hidden_dim))

  def forward(self, inputs, epsilon=1e-8):
    size = inputs.size()
    inputs = inputs.view(inputs.size(0), -1)
    sigma = t.std(inputs, -1)
    mean = t.mean(inputs, -1)
    output = (inputs - mean.unsqueeze(-1)) / (t.sqrt(sigma).unsqueeze(-1) + epsilon)
    output = self.alpha.repeat(size[0], size[1]) * output + self.beta.repeat(size[0], size[1])
    return output.view(size)


class PositionalEncoding(nn.Module):
  def __init__(self, embed_dim):
    super(PositionalEncoding, self).__init__()
    self.model_dim = embed_dim

  def forward(self, inputs):
    batch_size, max_len, _ = inputs.size()
    pe = t.zeros(max_len, self.model_dim)
    position = t.arange(1, max_len + 1, dtype=t.float32).unsqueeze(-1).repeat(1, self.model_dim // 2)
    div_term = t.pow(1e5, -t.arange(0, self.model_dim, 2,
                                    dtype=t.float32) / self.model_dim).unsqueeze(0).repeat(max_len, 1)
    pe[:, 0::2] = t.sin(position * div_term)
    pe[:, 1::2] = t.cos(position * div_term)
    pe = pe.unsqueeze(0).repeat(batch_size, 1, 1)
    return pe + inputs

Inference

Transformer最初是用来解决机器翻译(Machine Translation)的问题的。在训练阶段，给的是平行语料(源语言和目的语言)，但在训练完毕，使用或者测试时，并不会提供目的语言。这样Encoder的输入是什么呢？
一般来说，在原始序列的基础上首尾会分别加上<BOS>(Begin of Sequence)和<EOS>(End of Sequence)特殊字符。测试阶段，与RNN模型类似，Encoder端先输入<BOS>与Decoder端进行Attention操作，生成下一个字符，然后将生成的字符与Decoder输出进行Attention操作，以此类推，直到生成<EOS>或者序列长度大于预设最大长度为止。

2020-08-07: 最近使用Transformer进行seq2seq序列生成任务时遇到了一个问题，再Decoder端，不管输入是什么， 对Encoder输出的attention都是一样的。例如，Encoder端输入为"How are you ?"， Decoder端输入为"<BOS> I am"，目标输出为"fine"，查看中间结果发现Decoder输入的三个字符(<BOS>, I, am)对Encoder输出向量$\mathbf{e}\in {\mathbb{R}}^{l_e\times d_h}{l}_e=4$，进行attention的输出$\mathbf{a}\in {\mathbb{R}}^{l_d\times d_h}$几乎是一样的， i.e. $cosine\_similarity(\mathbf{a}[i],\mathbf{a}[j])\approx 1$，这说明了，输出的每个字对输入端的attention是无差别的，显然与我们的目标不一致。
再进行了两天的试错之后，发现有几个操作会增加向量的相似度（用余弦相似度衡量）。

1. softmax操作。在attention求权重和decoder的输出端，都会使用softmax将权重归一化。softmax会将向量/矩阵中所有元素压缩到[0,1]区间，当元素在0附近徘徊时，softmax的会将元素间的差别缩小，而transformer在进行softmax之前，会进行scale操作，将所有元素缩放至原始的$\frac{1}{\sqrt{d_m}}$，使得本来就很小的元素(使用xavier进行初始化，$\mathbf{w}\sim n(0,\frac{2}{\sqrt{n_{in}+n_{out}}})$)，变得更加趋近于零。输出元素会全部集中在$1/n_{out}$附近。从而使每个元素的注意力权重都相同。
2. Linear：这个比较玄学，个人查阅相关权重初始化的资料，发现自己的能力暂时无法证明这个操作带来的问题，但是做了一些模拟性的实验，发现这个操作确实能使不同权重的相关性增大。
3. 在去掉scaled-dot-multiplication中的scale操作甚至在原来的基础之上乘以$\sqrt{d_m}$之后，发现确实能够减小权重相关性(这也会使权重更新变慢，这也是transformer中进行$scale$的原因)。但在输出层，依然是输出一样的数据。原本代码的操作是这样的，Decoder端的输出字符个数与Decoder输入字符个数是相等的，也就是说输入(<BOS>, I, am)之后，会输出$dec\_output\in {\mathbb{R}}^{l_d\times vocab\_size}$，但我们要预测的是下一个字符，所以期望输出为$dec\_outpu{t}^{\prime }\in {\mathbb{R}}^{1\times vocab\_size}$，我使用最直接的办法，对原输出在第0维度(如果加上$batch$，那么就是第1维度)求均值，使得输出为长度为1。
   后来我重新思考了这种方法的合理性，将上下文进求均值得到next token貌似不是一个合理的解释。于是改成了Decoder中输出最后一个token(也就是$dec\_output[-1,:]$)作为最后输出，然后对其求argmax操作。实验发现，problem solved!
   在训练网络的时候，发现自己数学功底真的好弱。分析不出矩阵的性质(均值，方差等)对训练的影响，只能找文献、做实验试错，逐渐变成一个调参人员。