转载自:
原文链接:https://zhuanlan.zhihu.com/p/360343417
作者:Algernon
少量行文修改。

Transformer并没有特别复杂,但是理解Transformer对于初学者不是件容易的事,原因因在于Transformer的解读往往没有配套的简单的demo,并且缺少端到端的demo,就很难透彻理解Transformer的具体运算流程。Github上虽然有使用Transformer的翻译模型、推断模型,但作为demo来说,代码又太过复杂,不易上手。

本文将从以下三个部分以理论与实例相结合的方式阐述Transformer:

  1. 子模块解读:拆解Transformer,结合代码解读各个子模块的运算细节;
  2. 极简单翻译模型Demo:讲解使用transformer的翻译模型,将 ('<bos>', 'i', 'am', 'iron', 'man', '<eos>') 翻译为 ('<bos>', '我', '是', '钢铁', '侠', '<eos>') 的训练与推理过程。(训练与推理,都只翻译这一句话);
  3. Attention的mask作用:解读attention中mask的作用。

本文配套的源代码地址:https://github.com/thisiszhou/Transformer-Translate-Demo

1 子模块解读

1.1 MultiheadAttention

MultiheadAttention多头注意力,和注意力Attention稍有区别,是整个Transformer的核心,其他模块都是MultiheadAttention的封装与组合。下面先讲解Attention,暂时忽略batch_size,可以暂且理解其为batch_size等于1的特例。

Attention的计算流程如下图所示:

深度学习 – 以一个极简单的中英文翻译Demo彻底理解Transformer-StubbornHuang Blog

Attention的输入为QKV三个矩阵,其中tgt指target,src指source,代表的含义在不同的任务中有所差异。在本文的demo英译汉这个任务中,英语('<bos>', 'i', 'am', 'iron', 'man', '<eos>')就是src,汉语('<bos>', '我', '是', '钢铁', '侠', '<eos>')就是tgt,其中,tgt_size,src_size分别指汉语句子的最大长度以及英文句子的最大长度,emb_dim是每个单词词向量的维度(在其他任务中,emb_dim指的是特征维度)。

上述的Attention模型依次进行了如下操作:
1. 将QKV经过一层全连接层,得到新的Q^{\prime }K^{\prime }V^{\prime }
2. 将Q^{\prime }K^{\prime }的转置矩阵相乘,得到矩阵W,将W进行dim=-1维度的softmax操作,得到新的权重矩阵W^{\prime },其矩阵中的每一行i代表了tgt中的第i个词对src中每个词的注意力权重,所以W^{\prime }的维度为tgt_size*src_size;
3. 使用权重矩阵W^{\prime }与矩阵V^{\prime }相乘得到新的矩阵O,将O经过一层全连接层,得到输出Output;

以上为Attention的计算流程,Self-Attention,就是QKV输入为同一矩阵,即可计算矩阵关于自己的Attention。

现在我们对Self-Attention有了一定的了解,那么什么是MultiheadAttention呢?很多解读说是并行多个Attention,之后再合并起来,这种说法不完全对,MultiheadAttention是对emb_dim维度进行切分,之后并行Attention,精髓就在于对emb_dim维度进行切分。

MultiheadAttention计算流程如下图所示:

深度学习 – 以一个极简单的中英文翻译Demo彻底理解Transformer-StubbornHuang Blog

这里同样忽略batch_size的维度,计算步骤如下:

  1. QKV经过一层全连接层,得到新的Q^{\prime }K^{\prime }V^{\prime }
  2. Q^{\prime }K^{\prime }V^{\prime }从emb_dim的维度进行切分,共切割num_heads个,这里要求emb_dim可以被num_heads整除,切分后的结果为Q^{1}Q^{2},...,Q^{n}K^{1}K^{2},...,K^{n}V^{1}V^{2},...,V^{n}
  3. Q^{1}K^{1}的转置矩阵相乘得到矩阵W^{1}Q^{2}K^{2}的转置矩阵相乘得到矩阵W^{2},再将得到的W^{1}W^{2}、...、W^{n}进行dim=-1维度的softmax操作得到新的权重矩阵W^{\prime 1}W^{\prime 2}、...、W^{\prime n}
  4. 使用W^{1}乘以矩阵V^{1},得到矩阵O^{1},使用W^{2}乘以矩阵V^{2},得到矩阵O^{2},...,使用W^{n}乘以矩阵V^{n},得到矩阵O^{n}
  5. O^{i}的形状为tgt_size * head_size,其中head_size为emb_dim除以num_heads的商,一共有num_heads个O^{i},将这些O^{i}在head_size的维度进行拼接得到矩阵O,形状为tgt_size * emb_dim;
  6. 矩阵O经过一层全连接,得到输出Output;

MultiheadAttention模块的代码为:

class MultiheadAttention(Module):
    def __init__(self,
                 word_emb_dim,
                 nheads,
                 dropout_prob=0.
                 ):
        super(MultiheadAttention, self).__init__()
        self.word_emb_dim = word_emb_dim
        self.num_heads = nheads
        self.dropout_prob = dropout_prob
        self.head_dim = word_emb_dim // nheads
        assert self.head_dim * nheads == self.word_emb_dim  # embed_dim must be divisible by num_heads

        self.q_in_proj = Linear(word_emb_dim, word_emb_dim)
        self.k_in_proj = Linear(word_emb_dim, word_emb_dim)
        self.v_in_proj = Linear(word_emb_dim, word_emb_dim)

        self.out_proj = Linear(word_emb_dim, word_emb_dim)

    def forward(self,
                query: Tensor,
                key: Tensor,
                value: Tensor,
                key_padding_mask: Optional[Tensor] = None,
                attn_mask: Optional[Tensor] = None):
        """
        :param query: Tensor, shape: [tgt_sequence_size, batch_size, word_emb_dim]
        :param key:   Tensor, shape: [src_sequence_size, batch_size, word_emb_dim]
        :param value: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim]
        :param key_padding_mask:  Tensor, shape: [batch_size, src_sequence_size]
        :param attn_mask: Tensor, shape: [tgt_sequence_size, src_sequence_size]
        :return: Tensor, shape: [tgt_sequence_size, batch_size, word_emb_dim]
        """

        # 获取query的shape,这里按照torch源码要求,按照tgt_sequence_size, batch_size, word_emb_dim顺序排列
        tgt_len, batch_size, word_emb_dim = query.size()
        num_heads = self.num_heads
        assert word_emb_dim == self.word_emb_dim
        head_dim = word_emb_dim // num_heads

        # 检查word_emb_dim是否可以被num_heads整除
        assert head_dim * num_heads == word_emb_dim
        scaling = float(head_dim) ** -0.5

        # 三个Q、K、V的全连接层
        q = self.q_in_proj(query)
        k = self.k_in_proj(key)
        v = self.v_in_proj(value)

        # 这里对Q进行一个统一常数放缩
        q = q * scaling

        # multihead运算技巧,将word_emb_dim切分为num_heads个head_dim,并且让num_heads与batch_size暂时使用同一维度
        # 切分word_emb_dim后将batch_size * num_heads转换至第0维,为三维矩阵的矩阵乘法(bmm)做准备
        q = q.contiguous().view(tgt_len, batch_size * num_heads, head_dim).transpose(0, 1)
        k = k.contiguous().view(-1, batch_size * num_heads, head_dim).transpose(0, 1)
        v = v.contiguous().view(-1, batch_size * num_heads, head_dim).transpose(0, 1)

        src_len = k.size(1)

        # Q、K进行bmm批次矩阵乘法,得到权重矩阵
        attn_output_weights = torch.bmm(q, k.transpose(1, 2))
        assert list(attn_output_weights.size()) == [batch_size * num_heads, tgt_len, src_len]

        if attn_mask is not None:
            if attn_mask.dtype == torch.bool:
                attn_output_weights.masked_fill_(attn_mask, float('-inf'))
            else:
                attn_output_weights += attn_mask

        if key_padding_mask is not None:
            attn_output_weights = attn_output_weights.view(batch_size, num_heads, tgt_len, src_len)
            attn_output_weights = attn_output_weights.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2),
                float('-inf'),
            )
            attn_output_weights = attn_output_weights.view(batch_size * num_heads, tgt_len, src_len)

        # 权重矩阵进行softmax,使得单行的权重和为1
        attn_output_weights = torch.softmax(attn_output_weights, dim=-1)
        attn_output_weights = torch.dropout(attn_output_weights, p=self.dropout_prob, train=self.training)

        # 权重矩阵与V矩阵进行bmm操作,得到输出
        attn_output = torch.bmm(attn_output_weights, v)
        assert list(attn_output.size()) == [batch_size * num_heads, tgt_len, head_dim]

        # 转换维度,将num_heads * head_dim reshape回word_emb_dim,并且将batch_size调回至第1维
        attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, batch_size, word_emb_dim)

        # 最后一层全连接层,得到最终输出
        attn_output = self.out_proj(attn_output)
        return attn_output

1.2 TransformerEncoder

Transformer主要有TransformerEncoder和TransformerDecoder组成,一个是编码,一个解码。编码时只对src进行编码,例如本文中的例子,将英文翻译为中文,那么只对src英文输入语句进行encode。

TransformerEncoder计算流程如下图所示(图中省略激活层Relu以及Dropout):

深度学习 – 以一个极简单的中英文翻译Demo彻底理解Transformer-StubbornHuang Blog

TransformerEncoder由6个(数量可自定义)TransformerEncoderLayer组成,其中单个TransformerEncoderLayer的计算流程为:

  1. MultiheadAttention,此处Query、Key、Value都是同一Input,所以也是Self-Attention,后接LayerNorm;
  2. 一个shape为 emb_dim, dim_feedforward 的全连接层;
  3. 一个shape为dim_feedforward, emb_dim的全连接层,后接LayerNorm,输出Output作为当前TransformerEncoderLayer的输出;

由于TransformerEncoderLayer的输入shape为src_size * emb_dim,输出shape也为src_size * emb_dim,所以TransformerEncoderLayer的输出可以直接喂给下一个TransformerEncoderLayer,重复六次之后,就得到了TransformerEncoder的输出。

TransformerEncoder的代码为:

class TransformerEncoderLayer(Module):

    def __init__(self, word_emb_dim, nhead, dim_feedforward=2048, dropout_prob=0.1):
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(word_emb_dim, nhead, dropout_prob=dropout_prob)

        self.linear1 = Linear(word_emb_dim, dim_feedforward)
        self.dropout = Dropout(dropout_prob)
        self.linear2 = Linear(dim_feedforward, word_emb_dim)

        self.norm1 = LayerNorm(word_emb_dim)
        self.norm2 = LayerNorm(word_emb_dim)
        self.dropout1 = Dropout(dropout_prob)
        self.dropout2 = Dropout(dropout_prob)

        self.activation = torch.relu

    def forward(self, src: Tensor,
                src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        """
        :param src: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim]
        :param src_mask: Tensor, shape: [src_sequence_size, src_sequence_size]
        :param src_key_padding_mask: Tensor, shape: [batch_size, src_sequence_size]
        :return: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim]
        """
        # self attention
        src2 = self.self_attn(src, src, src,
                              attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        # 两层全连接
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src


class TransformerEncoder(Module):

    __constants__ = ['norm']

    def __init__(self, encoder_layer, num_layers, norm):
        super(TransformerEncoder, self).__init__()
        # 将同一个encoder_layer进行deepcopy n次
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self,
                src: Tensor,
                mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        """
        :param src: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim]
        :param mask: Tensor, shape: [src_sequence_size, src_sequence_size]
        :param src_key_padding_mask: Tensor, shape: [batch_size, src_sequence_size]
        :return: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim]
        """
        output = src
        # 串行n个encoder_layer
        for mod in self.layers:
            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)

        output = self.norm(output)
        return output


def _get_clones(module, N):
    return ModuleList([copy.deepcopy(module) for _ in range(N)])

1.3 TransformerDecoder

TransformerDecoder计算流程如下图所示

深度学习 – 以一个极简单的中英文翻译Demo彻底理解Transformer-StubbornHuang Blog

从上图可以看出,TransformerDecoder有两个输入,一个输出是Encoder的输出memory,一个是tgt。

这里tgt读者可能会有问题,Decoder对Encoder的输出memory进行解码理所应当,但是tgt是哪来的?

首先,没有tgt输入,只对memory解码的Decoder也是存在的,但是在翻译任务中,需要有一个额外的tgt输入,来得到不同的输出。当前先讲解Decoder的运算流程,这里tgt具体应用,在第二章翻译demo中会详细解释。

TransformerDecoder也是由六个(数量可自定义)TransformerDecoderLayer串联组成,其中单个的TransformerDecoderLayer的计算流程如下:

  1. tgt进行self-attention,经过LayerNorm,当作MultiheadAttention的Query;
  2. memory当作MultiheadAttention的Key、Value,结合Query,进行一次MultiheadAttention,经过LayerNorm后输出;
  3. 经过一个shape为 emb_dim, dim_feedforward 的全连接层,和一个shape为dim_feedforward, emb_dim的全连接层,后接LayerNorm,输出tgt_out,shape与最开始的输入tgt相同。

TransformerEncoder的代码为:

class TransformerDecoderLayer(Module):

    def __init__(self, word_emb_dim, nhead, dim_feedforward=2048, dropout_prob=0.1):
        super(TransformerDecoderLayer, self).__init__()
        # 初始化基本层
        self.self_attn = MultiheadAttention(word_emb_dim, nhead, dropout_prob=dropout_prob)
        self.multihead_attn = MultiheadAttention(word_emb_dim, nhead, dropout_prob=dropout_prob)
        # Implementation of Feedforward model
        self.linear1 = Linear(word_emb_dim, dim_feedforward)
        self.dropout = Dropout(dropout_prob)
        self.linear2 = Linear(dim_feedforward, word_emb_dim)

        self.norm1 = LayerNorm(word_emb_dim)
        self.norm2 = LayerNorm(word_emb_dim)
        self.norm3 = LayerNorm(word_emb_dim)
        self.dropout1 = Dropout(dropout_prob)
        self.dropout2 = Dropout(dropout_prob)
        self.dropout3 = Dropout(dropout_prob)

        self.activation = torch.relu

    def forward(self,
                tgt: Tensor,
                memory: Tensor,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        """
        :param tgt:                     Tensor, shape: [tgt_sequence_size, batch_size, word_emb_dim]
        :param memory:                  Tensor, shape: [src_sequence_size, batch_size, word_emb_dim]
        :param tgt_mask:                Tensor, shape: [tgt_sequence_size, tgt_sequence_size]
        :param memory_mask:             Tensor, shape: [src_sequence_size, src_sequence_size]
        :param tgt_key_padding_mask:    Tensor, shape: [batch_size, tgt_sequence_size]
        :param memory_key_padding_mask: Tensor, shape: [batch_size, src_sequence_size]
        :return:                        Tensor, shape: [tgt_sequence_size, batch_size, word_emb_dim]
        """
        # tgt的self attention
        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        # tgt与memory的attention
        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        # 两层全连接层
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt


class TransformerDecoder(Module):
    __constants__ = ['norm']

    def __init__(self, decoder_layer, num_layers, norm):
        super(TransformerDecoder, self).__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        """
        :param tgt:                     Tensor, shape: [tgt_sequence_size, batch_size, word_emb_dim]
        :param memory:                  Tensor, shape: [src_sequence_size, batch_size, word_emb_dim]
        :param tgt_mask:                Tensor, shape: [tgt_sequence_size, tgt_sequence_size]
        :param memory_mask:             Tensor, shape: [src_sequence_size, src_sequence_size]
        :param tgt_key_padding_mask:    Tensor, shape: [batch_size, tgt_sequence_size]
        :param memory_key_padding_mask: Tensor, shape: [batch_size, src_sequence_size]
        :return:                        Tensor, shape: [tgt_sequence_size, batch_size, word_emb_dim]
        """
        output = tgt

        for mod in self.layers:
            output = mod(output, memory, tgt_mask=tgt_mask,
                         memory_mask=memory_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=memory_key_padding_mask)

        output = self.norm(output)

        return output


def _get_clones(module, N):
    return ModuleList([copy.deepcopy(module) for _ in range(N)])

1.4 PositionalEncoding

PositionalEncoding虽然在Transformer中是第一个模块,但是这里最后讲解,因为PositionalEncoding对于Transformer来说,不是必须的,在torch.nn.Transformer中,没有包含PositionalEncoding,需要自己实现。有一些序列化含义不强的场景,PositionalEncoding可以省略。

在翻译任务中,当英文句子被表示成一个矩阵(每一行是句子中对应位置英文单词的词向量),位置信息被淡化,所以PositionalEncoding的作用,就是体现每个词的相对位置与绝对位置信息。

论文中的位置编码如下:

\begin{array}{c}
P E(p o s, 2 i)=\sin \left(p o s / 10000^{2 i / d_{\text {model }}}\right) \\
P E(p o s, 2 i+1)=\cos \left(p o s / 10000^{2 i / d_{\text {model }}}\right)
\end{array}

PositionalEncoding并不是对句子的位置进行一维编码,而是对句子的位置position、每个单词的词向量emb_dim,共两个维度进行编码,而且是独立编码,两个维度互不干涉。上述公式中,pos就是当前词在句子的位置,2i2i+1,是词向量emb_dim中的位置,所以PositionalEncoding编码矩阵的形状为 (tgt_sequence_size, word_emb_dim)。

PositionalEncoding的代码为:

class PositionalEncoding(nn.Module):

    def __init__(self, word_emb_dim: int, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        position_emb = torch.zeros(max_len, word_emb_dim)

        # position 编码
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

        dim_div_term = torch.exp(torch.arange(0, word_emb_dim, 2).float() * (-math.log(10000.0) / word_emb_dim))

        # word_emb_dim 编码
        position_emb[:, 0::2] = torch.sin(position * dim_div_term)
        position_emb[:, 1::2] = torch.cos(position * dim_div_term)
        pe = position_emb.unsqueeze(0).transpose(0, 1)  # shape: (max_len, 1, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x: Tensor):
        """
        :param x: Tensor, shape: [batch_size, sequence_length, word_emb_dim]
        :return: Tensor, shape: [batch_size, sequence_length, word_emb_dim]
        """

        # 编码信息与原始信息加和后输出
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

PositionalEncoding位置编码,与原始信息加和后输出。

2 极简的翻译模型Demo

Transformer的结构并不是固定的,上面所述的四种基础结构(MultiheadAttention,TransformerEncoder,TransformerDecoder,PositionalEncoding)基本是固定的,而Transformer可以由自由组合,torch.nn.Transformer官方demo中,Decoder就是一个全连接层,而没有用上述TransformerDecoder。

2.1 明确翻译任务

准备示例数据。

假设我们已经做好了英文词典和中文词典,并且对每一个字符编号:

cn_dict = {
    '<bos>': 0,
    '<eos>': 1,
    '<pad>': 2,
    '是': 3,
    '千': 4,
    '你': 5,
    '万': 6,
    '在': 7,
    '我': 8,
    '人': 9,
    '三': 10,
    '一': 11,
    '侠': 12,
    '遍': 13,
    '二': 14,
    '爱': 15,
    '好': 16,
    '钢铁': 17
}

en_dict = {
    '<bos>': 0,
    '<eos>': 1,
    '<pad>': 2,
    'i': 3,
    'three': 4,
    'am': 5,
    'love': 6,
    'you': 7,
    'he': 8,
    'times': 9,
    'is': 10,
    'thousand': 11,
    'hello': 12,
    'iron': 13,
    'man': 14
}

语料也已经准备好了,只有如下两句话,每个单词都被收录在上述词典中:

sentence_pair_demo = [
    [
        ('<bos>', 'i', 'am', 'iron', 'man', '<eos>'),
        ('<bos>', '我', '是', '钢铁', '侠', '<eos>')
    ],
    [
        ('<bos>', 'i', 'love', 'you', 'three', 'thousand', 'times', '<eos>'),
        ('<bos>', '我', '爱', '你', '三', '千', '遍', '<eos>')
    ]
]

本文翻译模型demo,只训练翻译两句话。了解Transformer最简易翻译模型的原理,只包含两句话的数据集足够。真正使用时,只需要替换更大的训练数据集即可。

注:当前例句假设已经分好词,并且一句话有完整的开始标记符'<bos>',和结束标记符'<eos>'

2.2 基于Transformer的翻译模型结构

深度学习 – 以一个极简单的中英文翻译Demo彻底理解Transformer-StubbornHuang Blog

Transformer的步骤为:

  1. 将tgt和src经过word_emb得到各个词的词向量,并且经过PositionalEncoding,得到含有position信息的src tensor和tgt tensor;
  2. 将src经过TransformerEncoder得到memory;
  3. 将tgt与memory送入TransformerDecoder得到tgt_out;
  4. 最后经过reshape与一层全连接层,得到一个长度为tgt词典长度的向量,代表每个词此时的预测概率(需要经过softmax)

注意src与tgt通过<pad>填充为固定长度,整个模型的设计为每次输入翻译前原句和翻译后已经获得的词,来预测下一个词,例如:输入为 src = ('<bos>', 'i', 'am', 'iron', 'man', '<eos>', '<pad>', '<pad>')tgt = ('<bos>', '我', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>'),那么模型下一个预测的词,应该为 '是'。

2.3 模型的训练

就算是只有两句话的数据集,也需要有一个Dataset工具来组织batch数据(详见代码utils.dataset)。

这里Dataset的实现不展开,和Transformer关系不大,只需要明确一下get_batch函数的输出:

def get_batch(self, batch_size=2, padding_str='<pad>', need_padding_mask=False):
    """
    :return: src, tgt_in, tgt_out, src_padding_mask, tgt_padding_mask
    src:              Tensor, shape: [batch_size, src_sequence_size]
    tgt_in:           Tensor, shape: [batch_size, tgt_sequence_size]
    tgt_out:          Tensor, shape: [batch_size, tgt_sequence_size]
    src_padding_mask: Tensor, shape: [batch_size, src_sequence_size]
    tgt_padding_mask: Tensor, shape: [batch_size, tgt_sequence_size]
    """

同样,这里暂时忽略mask。在训练的时候,[('<bos>', 'i', 'am', 'iron', 'man', '<eos>'), ('<bos>', '我', '是', '钢铁', '侠', '<eos>')],是一组数据,单batch的一些数据示例如下:

1) src = ('<bos>', 'i', 'am', 'iron', 'man', '<eos>', '<pad>', '<pad>'),tgt = ('<bos>', '我', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>'),tgt_out = '是'
2) src = ('<bos>', 'i', 'am', 'iron', 'man', '<eos>', '<pad>', '<pad>'),tgt = ('<bos>', '我', '是', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>'),tgt_out = '钢铁'
3) src = ('<bos>', 'i', 'am', 'iron', 'man', '<eos>', '<pad>', '<pad>'),tgt = ('<bos>', '我', '是', '钢铁', '<pad>', '<pad>', '<pad>', '<pad>'),tgt_out = '侠'

如上所示,当前翻译Demo设计的是,输入英文整句,以及中文已经翻译出的词,来预测下一个词。

例如('<pad>'为填充字符,旨在将每句话的长度拉齐):

1)第一步:输入为src = ('<bos>', 'i', 'am', 'iron', 'man', '<eos>', '<pad>', '<pad>')tgt = ('<bos>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>'),期望输出为 '我';
2)第二步:输入为src = ('<bos>', 'i', 'am', 'iron', 'man', '<eos>', '<pad>', '<pad>')tgt = ('<bos>', '我', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>'),期望输出为'是';

重复上述步骤,直到整句话输出为'<eos>'或者达到最大长度后停止。

那么在Transformer的训练中,输入是src = ('<bos>', 'i', 'am', 'iron', 'man', '<eos>', '<pad>', '<pad>'),tgt = ('<bos>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>'),输出的标签就是'我'。当前这里所有的词,都会转换成该词在词典中的序号,并且按照batch拼成张量。一句话的长度是src_sequence_size或者tgt_sequence_size,按照batch输出后,Tensor的shape为[batch_size, src_sequence_size]或者[batch_size, tgt_sequence_size]

2.4 预测

预测最主要的函数如下

def infer_with_transformer(model: Transformer, src: Tensor, tgt_dict: Dictionary, max_length=8) -> List[str]:
    out_seq = ['<bos>']
    predict_word = ''
    while len(out_seq) < max_length and predict_word != '<eos>':
        tgt_in = transform_words_to_tensor(out_seq, tgt_dict)
        output = model(src, tgt_in)
        word_i = torch.argmax(output, -1).item()
        predict_word = tgt_dict.i2w[word_i]
        out_seq.append(predict_word)
    return out_seq

在最开始,out_seq只包含一个'<bos>',每次预测出的词,都append进out_seq,再循环预测下一个词,直到输出为'<eos>'或者达到最大长度后停止。

使用训练好的模型,可以看到以下输出:

Input sentence: ['<bos>', 'i', 'love', 'you', 'three', 'thousand', 'times', '<eos>']
After translate: ['<bos>', '我', '爱', '你', '三', '千', '遍', '<eos>']
Input sentence: ['<bos>', 'i', 'am', 'iron', 'man', '<eos>']
After translate: ['<bos>', '我', '是', '钢铁', '侠', '<eos>']

3 Attention的mask作用

在MultiheadAttention中,一共有两个mask,一个是key_padding_mask,一个是attn_mask。两个mask相互是独立的,和Multihead没有太大关系,所以用Attention来讲解mask的作用。

回顾以下Attention的计算流程,

深度学习 – 以一个极简单的中英文翻译Demo彻底理解Transformer-StubbornHuang Blog

Query和Key的作用,是计算每一个tgt关于每一个src的权重,所以W矩阵(注意是在softmax前的W矩阵)的shape为tgt_size * sec_size,两个mask的作用就是体现在这里。

假设当前W矩阵如下(这里用1填充,方便演示计算):

深度学习 – 以一个极简单的中英文翻译Demo彻底理解Transformer-StubbornHuang Blog

3.1 key_padding_mask

注意到W中,src有填充符号<pad>,该字符并不需要被tgt注意到,因为对于翻译没有作用。而W矩阵中,每一个tgt字符,都有关于<pad>字符的注意力权重,这里只需要一个key_padding_mask矩阵,进行如下操作,即可消除tgt对于src中<pad>的注意力权重:

深度学习 – 以一个极简单的中英文翻译Demo彻底理解Transformer-StubbornHuang Blog

在MultiheadAttention的forward中,key_padding_mask参数传入的是一个Tensor,shape为[batch_size, src_sequence_size],可以为bool型(padding的位置为True),也可以为float型(padding位置为-inf,其余为0)。

3.2 attn_mask

attn_mask与key_padding_mask相互独立,其shape为[tgt_sequence_size, src_sequence_size],旨在控制tgt关于src的注意力权重。

例如,比较常见的是上三角attn_mask矩阵:
深度学习 – 以一个极简单的中英文翻译Demo彻底理解Transformer-StubbornHuang Blog

加入上三角attn_mask矩阵后,tgt的第一个词,只能关注到src的第一个词,tgt的第二个词,只能关注到src的第一个词和第二个词,... 。

本文中的demo,并不需要加入attn_mask,因为在预测第二个词时,神经网络的输入只有之前的词。而在有些任务中,前序列元素不需要关注后序列元素,就可以使用上三角attn_mask来控制。