深度学习 Transformer 架构详解:代码 + 图示

本文通过将图片和代码结合的方式,讲解Transformer的架构。主要内容来源于:The Illustrated TransformerThe Annotated Transformer

本文内容

  1. 从宏观的角度看
    1. 编码器(Encoder)和解码器(Decoder)
    2. 编码器层(Encoder Layer)
    3. 解码器层(Decoder Layer)
  2. 注意力机制(Attention)
    1. 自注意力(Self-Attention)
    2. 自注意力的计算(单个)
    3. 自注意力的计算(矩阵)
    4. 注意力机制(输入输出)
    5. 多头注意力(代码)
    6. 多头注意力
  3. 编码器(Encoder)
    1. 使用位置编码表示序列的位置
    2. 全连接的前馈网络(Feed-Forward Networks)
    3. 子层之间的连接(残差和层归一化)
      1. 原始论文
      2. 参考代码
  4. 解码器(Decoder)
  5. 最后的线性层和softmax层
  6. 嵌入向量(Embeddings)
  7. 训练过程
  8. 损失函数
  9. 整个模型

Transformer模型是一个利用注意力机制来提高模型训练速度的先进模型。该模型是由论文 Attention is All You Need (2017.06,Google) 所提出的。为了方便理解,我们将逐一深入介绍其各个概念。

注意(由于本文参考了不同的文档,部分称呼不一致,这里统一说明一下):

  1. 编码器:在论文中,编码器被称之为“Encoder”,在代码中为“Encoder”类,而在图示中则表述为“Encoders”。
  2. 编码器层:在论文中,编码器层被表述为“encoder layer”,在代码中为“EncoderLayer”类,而在图示中则表述为“Encoder”。
  3. 类似地,解码器也有相应的称呼区别。
  4. 在中文翻译中,我们采用论文中的称呼。
  5. 在后续的示例中,我们假设输入的句子的长度为 S 个 token,模型参数 D(即 dmodel)为 512,多头注意力机制的头数 h=8,多头注意力权重矩阵对应的维度 d = 512/8 = 64。

从宏观的角度看

我们可以将整个Transformer模型视为一个黑盒,在机器翻译任务中,它可以将句子从一种语言翻译成另一种语言。

编码器(Encoder)和解码器(Decoder)

打开这个黑盒,我们首先可以看到一个编码器(Encoder)和一个解码器(Decoder),以及二者之间存在某种关联。

进一步探索,编码器由 N=6N=6 个相同的编码器层(encoder layer)堆叠而成,同样,解码器也由 6 个相同的层堆叠而成。6 是论文中的一个设定,可以换成其他值。

# 本示例 pytorch 代码仅包含模型架构和示例推理,未包含模型训练。可以拷贝后执行。
import torch
import torch.nn as nn
import math
import copy

def clones(module, N):
    """生成 N 个相同的层"""
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class Encoder(nn.Module):

    def __init__(self, layer, N):
        """
        编码器由 N 层组成
        layer: EncoderLayer,详见后
        """
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size) 

    def forward(self, x, mask):
        """将输入(input)和掩码(mask)顺序通过每一个子层,掩码不变。"""
        for layer in self.layers:
            x = layer(x, mask)
        # 代码 SublayerConnection(见后)是先做 LayerNorm,最后做残差。所以这里最后归一化 1 次。
        return self.norm(x)  
# Decoder 也是由 6 个相同的层组成
class Decoder(nn.Module):
    
    # layer: DecoderLayer,详见后
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    # memory: 来自self_attn
    def forward(self, x, memory, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

6 个编码器层的结构是相同的(但是他们之间的权重是不共享的),每个编码器层都可以分为 2 个子层。

编码器层(Encoder Layer)

编码器层的输入首先进入自注意力子层(Self-Attention),该子层的作用在于帮助编码器关注句子中的其他词汇,以便更好地编码某个特定词汇。

随后,自注意力子层的输出将传递给一个前馈神经网络(Feed-Forward Neural Network)。结构完全相同的前馈网络被独立地应用于每个位置。

输入输出对理解数据流非常重要。编码器层的输入形状为 S x D(请参见下面的图表),其中 S 是源句子长度(例如,英语句子),而 D 是嵌入的维度(也是模型维度,论文中取值为 512)。

编码器的输入和输出形状相同。由于编码器层是相互叠加的,因此,我们希望其输出具有与输入相同的维度,以便它可以轻松地流入下一个编码器层。因此,输出也是 S x D 形状。

编码器层的输入和输出是相同的

class SublayerConnection(nn.Module):
    """
    该类为子层(sublayer)之间通用的连接:残差连接(A residual connection)+ 层归一化(A layer norm)+ Dropout
    图示见后面的"子层之间的连接(残差和层归一化)"章节
    子层的类别有(见后面的示例代码):自注意力层(MultiHeadedAttention)、全连接层(PositionwiseFeedForward)
    顺序执行:norm, sublayer, dropout, resnet
    参考 EncoderLayer.forward 中的示例调用顺序。
    """
    
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        # 这里和原始论文有出入,做的 Pre_LN,详见后面的说明。
        return x + self.dropout(sublayer(self.norm(x)))
class EncoderLayer(nn.Module):
    """
    编码器层由 self-attn、feed_forward 和子层连接(SublayerConnection)组成。
    self_attn: MultiHeadedAttention,详见后面代码。
    feed_forward: PositionwiseFeedForward,详见后面代码。
    """
    
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        # 每个层都有两个子层。
        # 第一个是multi-head self-attention
        self.self_attn = self_attn
        # 第二个是简单的全连接前馈网络。
        self.feed_forward = feed_forward
        # 子层之间的连接,有 2 个
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size

    def forward(self, x, mask):
        """顺次执行:norm, self_attn, dropout, resnet; norm, feed_forward, dropout, resnet"""
        # 编码器层的 x 参数同时用作 q,k,v,计算self_attn
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        # 全连接前馈网络
        return self.sublayer[1](x, self.feed_forward)

解码器层(Decoder Layer)

解码器层(decoder layer)也包含前面编码器中提到的两个层,不过区别在于这两个层之间还夹了一个注意力层(Encoder-Decoder Attention)。这个额外的注意力层的作用在于让解码器能够注意到输入句子中与解码任务相关的部分。

class DecoderLayer(nn.Module):
    """Decoder is made of self-attn, src-attn, and feed forward (defined below)"""

    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        """
        self_attn: MultiHeadedAttention,详见后
        src_attn: MultiHeadedAttention,详见后
        feed_forward: PositionwiseFeedForward,详见后
        """
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn

        # 解码器层相对编码器层多的一个子层(Encoder-Decoder Attention)
        self.src_attn = src_attn
        
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, src_mask, tgt_mask):
        """论文中图1右"""
        m = memory
        # 第1次attn
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        # 第2次attn
        # key, value矩阵来自Encoder的输出,query来自上一层DecoderLayer
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        # 第3子层为feed_forward
        return self.sublayer[2](x, self.feed_forward)

在一个已经训练好的Transformer模型中,输入是怎么变为输出的呢?首先我们要知道各种各样的张量(向量)是如何在这些组件之间变化的。

与其他的NLP项目一样,我们首先需要把输入的每个单词通过词嵌入(embedding)转化为对应的向量。 每个单词被嵌入到大小为 512 的向量(图中用 4 格方框表示)

所有编码器层接收一组向量作为输入(论文中的输入向量的维度是512)。最底下的那个编码器层接收的是嵌入向量,之后的编码器层接收的是前一个编码器层的输出。

向量列表的长度这个超参数是我们可以设置的,一般来说是我们训练集中最长的那个句子的长度。

当我们的输入序列经过词嵌入之后得到的向量会依次通过编码器层中的两个层。

参考EncoderLayer.forward代码,包含句子中所有token的整个 X 矩阵,作为参数,传入自注意力层,分别生成Q,K,V,然后进行运算。

接下来,我们用一个短句(Thinking Machine)作为例子,看看在编码器层的每个子层中发生了什么。

注意力机制(Attention)

注意力机制是论文的核心,它在编码器和解码器部分的处理稍有差异。让我们先以编码器部分的注意力层机制为例进行介绍。

上边我们已经说了,每个编码器层接受一组向量作为输入。在其内部,输入向量先通过一个自注意力层,再经过一个前馈神经网络层,最后将其将输出给下一个编码器层。

不同位置上的单词都要经过自注意力层的处理,之后都会经过一个完全相同的前馈神经网络。

在这里,我们开始看到 Transformer 的一个关键特点,即每个位置上的单词在编码器层中有各自的流通方向。

  1. 在自注意力层中,这些路径之间存在依赖关系。单词和单词之间会有关联,假设一个句子有 50 个单词,那么可以粗略想象成自注意力计算过程中,会构造一个 50 x 50 的关联矩阵。
  2. 前馈神经网络(Feed Forward)层中没有这些依赖关系。每个单词独立通过前馈神经网络,单词和单词之间没有关联,因此各种路径可以在流过前馈网络层的时候并行计算。

自注意力(Self-Attention)

现在让我们看一下自注意力机制。

假设我们要翻译下边这句话:
”The animal didn't cross the street because it was too tired”

这里it指的是什么?是street还是animal?人理解起来很容易,但是对算法来讲就不那么容易了。

当模型处理it这个词的时候,自注意力会让itanimal关联起来。

当模型编码每个位置上的单词的时候,自注意力的作用就是:看一看输入句子中其他位置的单词,试图寻找一种对当前单词更好的编码方式。

如果你熟悉 RNNs 模型,回想一下 RNN 如何处理当前时间步的隐藏状态:将之前的隐藏状态与当前位置的输入结合起来。在 Transformer 中,自注意力机制也可以将其他相关单词的“理解”融入到我们当前处理的单词中。

当我们在最后一个encoder组建中对it进行编码的时候,注意力机制会更关注The animal,并将其融入到it的编码中。

自注意力的计算(单个)

先画图用向量解释一下自注意力是怎么算的,之后再看一下实际实现中是怎么用矩阵算的。

第一步 对于编码器的每个输入向量x,都会计算三个向量,即query、key和value向量。

这些向量的计算方法是将输入的词嵌入向量与三个权重矩阵相乘。这些权重矩阵是在模型训练阶段通过训练得到的。

什么是 “query”、“key”、“value” 向量?这三个向量是计算注意力时的抽象概念,请继续往下看注意力计算过程。

第二步 计算注意力得分。

假设我们现在在计算输入中第一个单词 Thinking 的自注意力。我们需要使用自注意力给输入句子中的每个单词打分,这个分数决定当我们编码某个位置的单词的时候,应该对其他位置上的单词给予多少关注度。

这个得分是query和key的点乘积得出来的。例如,如果我们处理位置#1的单词的自我注意,第一个分数将是q1和k1的点积。第二个分数是q1和k2的点积。(备注:在使用矩阵处理时,是用 Q 和 K 的转置相乘得到,详见后)。

第三步 将计算获得的注意力分数除以 8。

为什么选 8?是因为key向量的维度是 64,取其平方根,这样让梯度计算的时候更稳定。默认是这么设置的,当然也可以用其他值。

第四步 除 8 之后将结果扔进 softmax 计算,使结果归一化,softmax 之后注意力分数相加等于 1,并且都是正数。

这个 softmax 之后的注意力分数表示 在计算当前位置的时候,其他单词受到的关注度的大小。显然在当前位置的单词肯定有一个高分,但是有时候也会注意到与当前单词相关的其他词汇。

第五步 将每个 value 向量乘以注意力分数。这是为了强化我们想要关注的单词的 value,并尽量抑制其他不相关的单词(通过乘以一个接近于零的数,如 0.001)。这个过程被称为“缩放”或者“加权”,可以使得我们更加关注与目标单词相关的单词。

第六步 将上一步的结果相加,输出本位置的注意力结果。

这就是自注意力的计算。计算得到的向量直接传递给前馈神经网络。但是为了处理的更迅速,实际是用矩阵进行计算的。接下来我们看一下怎么用矩阵计算。

自注意力的计算(矩阵)

第一步是计算 Query、Key 和 Value 矩阵。我们通过将嵌入打包到矩阵 X 中,并将其乘以我们训练的权重矩阵(WQWKWVW^Q、W^K、W^V)来实现这一点。

X矩阵中的每一行对应于输入句子中的一个单词。
可再次看到嵌入向量维度(512,图中的 4 个框)和 q/k/v 向量维度(64,图中的 3 个框)的差异

最后,由于我们使用矩阵计算,因此可以将步骤 2 到 6 合并为一个公式,以计算自注意力层的输出。

几点说明:

  1. 如下公式,对于编码器,Q、K、V 都是输入矩阵 X;对于解码器的某些自注意力,其 K、V 来源于编码器的输出,详见后面的解释。
  2. 输入参数和权重矩阵点乘后的结果(QWiQ,KWiK,VWiVQW^Q_i, KW^K_i, VW^V_i),会再作为 Attention 函数的参数(Q,K,V),从而计算最终的值,详见 MultiHeadedAttention 类和 attention 函数。
MultiHead(Q,K,V)=Concat(head1,...,headh)WOheadi=Attention(QWiQ,KWiK,VWiV)\mathrm{MultiHead}(Q, K, V) = \\ \mathrm{Concat}(\mathrm{head_1}, ..., \mathrm{head_h})W^O \\ \mathrm{head_i} = \mathrm{Attention}(QW^Q_i, KW^K_i, VW^V_i)

论文中并行注意力层数量(简称头)为 h=8h=8 ,参数矩阵 WQ,WK,WVW^Q, W^K, W^V 的维度都是 dk=dv=dmodel/h=64d_k=d_v=d_{\text{model}}/h=64,即为512×64512 \times 64的矩阵( WiQRdmodel×dkW^Q_i \in \mathbb{R}^{d_{\text{model}} \times d_k}WiKRdmodel×dkW^K_i \in \mathbb{R}^{d_{\text{model}} \times d_k}WiVRdmodel×dvW^V_i \in \mathbb{R}^{d_{\text{model}} \times d_v})。 它们是在模型训练阶段学习得到的。

连接参数矩阵为512×512512 \times 512WORhdv×dmodelW^O \in \mathbb{R}^{hd_v \times d_{\text{model}}})。

通过将输入向量 x 与注意力头的权重矩阵相乘,可以得到对应的 query、key 和 value 向量。单个头获取的这三个向量维度是64,比嵌入向量的维度小,8个头的输出连接后变为 512。因此嵌入向量、编码器层的输入输出维度都是512。

将x1和WQ权重矩阵相乘得到q1(该单词的“query”向量)。按这样的方法,最终我们给输入的每一个单词都计算出一个“query”、一个 “key”和一个 “value”。

注意力机制(输入输出)

我们再结合输入输出的形状加深一下理解。

对于上图的解释:

  1. 假定输入的英文句子是“The quick brown fox“,句子长度 S 为4,参考“编码器层”章节的解释,注意力子层的输入形状为(4 x 512)。
  2. 自注意力层使用三个权重矩阵进行初始化——Query(WqW_q)、Key(WkW_k)和Value(WvW_v)。这些权重矩阵的尺寸都是 D x d,在论文中d取值为64,即权重矩阵的尺寸为 512 x 64。在训练模型时,我们将训练这些矩阵的权重。
  3. 在第一次计算(图中的Calc 1)中,我们通过将输入(注意:代码实现中是三个不同的输入,编码器层都是X,解码器层不同,见代码中的解释)与各自的Query、Key和Value权重矩阵相乘,计算出Q、K和V矩阵(尺寸为 S x d,示例中为 4 x 64)。
  4. 在第二次计算中,参考Attention计算公式,首先将Q和Kᵀ矩阵相乘,得到一个尺寸为 S x S(示例中为 4 x 4)的矩阵,然后将其除以√d的标量。然后对矩阵进行softmax运算,使得每一行的和都为1。这个矩阵可以理解为句子中每个词之间的关联度。
  5. 上面 S x S 的矩阵再和V矩阵相乘,得到尺寸为 S x d(示例中为 4 x 64)的矩阵。经过后续的连接操作后,传入下一层。

多头注意力(代码)

由于使用矩阵运算和pytorch库,代码实现非常简洁,请参考注意力机制和多头注意力章节理解。

论文中 Attention 计算的公式为:

Attention(Q,K,V)=softmax(QKTdk)V\mathrm{Attention}(Q, K, V) = \mathrm{softmax}(\frac{QK^T}{\sqrt{d_k}})V

如下是该公式的代码实现。

# 计算论文中图2左的'Scaled Dot Product Attention'
def attention(query, key, value, mask=None, dropout=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

MHA的整个流程实现如下。

class MultiHeadedAttention(nn.Module):
    # h即头的数量,论文中取8
    # d_model 即输入向量的维度,论文中取512。
    def __init__(self, h, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0

        # K矩阵维度数为512/8=64
        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h
        
        # 4 个大参数矩阵(512x512),对应投影权重矩阵 Wq, Wk, Wv 和连接矩阵 O
        # 参考论文图2
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    # 参数query,key, value的维度都为(S,512),来源如下:
    #     encoder:参数query, key, value为上一层的输出X
    #     decoder: self_attn时,参数query, key, value为X矩阵(output或者上一层DecoderLayer输出)
    #              src_attn时,参数key, value矩阵来自Encoder的输出矩阵(叫memory),query为上一层DecoderLayer输出 X
    # 
    # 中间计算的也出现了临时的query, key, value矩阵名称,注意区分
    #
    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k
        query, key, value = [
            lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears, (query, key, value))
        ]

        # 2) Apply attention on all the projected vectors in batch.
        x, self.attn = attention(
            query, key, value, mask=mask, dropout=self.dropout
        )

        # 3) "Concat" using a view and apply a final linear.
        x = (
            x.transpose(1, 2)
            .contiguous()
            .view(nbatches, -1, self.h * self.d_k)
        )
        del query
        del key
        del value
        return self.linears[-1](x)

多头注意力

论文进一步改进了自注意力层,增加了一个机制,也就是多头注意力机制。这样做有两个好处:

第一个好处,它扩展了模型专注于不同位置的能力。

在上面例子里只计算一个自注意力的的例子中,编码“Thinking”的时候,虽然最后 Z1 或多或少包含了其他位置单词的信息,但是它实际编码中还是被“Thinking”单词本身所支配。

如果我们翻译一个句子,比如“The animal didn’t cross the street because it was too tired”,我们会想知道“it”指的是哪个词,这时模型的“多头”注意力机制会起到作用。

第二个好处,它给了注意层多个“表示子空间”。

就是在多头注意力中同时用多个不同的 WVW^V 权重矩阵(Transformer 使用8个头部,因此我们最终会得到8个计算结果),每个权重都是随机初始化的。经过训练每个 WVW^V 都能将输入的矩阵投影到不同的表示子空间。

多头注意力的每个头有独立的Q/K/V权重矩阵(记为Wq,Wk,Wv),将 X 乘以这些矩阵,产生了不同的 Q/K/V 矩阵值。

如果我们做和上面相同的自注意力计算,只不过八次使用不同的权重矩阵,我们最后得到八个不同的Z矩阵

但是这会存在一点问题,多头注意力出来的结果会进入一个前馈神经网络,这个前馈神经网络可不能一下接收8个注意力矩阵,它的输入需要是单个矩阵(矩阵中每个行向量对应一个单词),所以我们需要一种方法把这8个压缩成一个矩阵。

怎么做呢?我们将这些矩阵连接起来,然后将乘以一个附加的权重矩阵

以上就是多头自注意力的全部内容。让我们把多头注意力上述内容 放到一张图里看一下子:

现在我们已经看过什么是多头注意力了,让我们回顾一下之前的一个例子,再看一下编码“it”的时候每个头的关注点都在哪里:

编码it,用两个head的时候:其中一个更关注the animal,另一个更关注tired。此时该模型对it的编码。除了it本身的表达之外,同时也包含了the animal和tired的相关信息

如果我们把所有的头的注意力都可视化一下,就是下图这样,但是看起来事情好像突然又复杂了。

编码器(Encoder)

除了注意力子层,其他模型子层说明如下。

使用位置编码表示序列的位置

其他参考:Transforme 结构:位置编码详解

到现在我们还没提到过如何表示输入序列中词汇的位置。

Transformer 在每个输入的嵌入向量中添加了位置向量。这些位置向量遵循某些特定的模式,这有助于模型确定每个单词的位置或不同单词之间的距离。将这些值添加到嵌入矩阵中,一旦它们被投射到Q、K、V中,就可以在计算点积注意力时提供有意义的距离信息。

为了让模型能知道单词的顺序,我们添加了位置编码,位置编码是遵循某些特定模式的。

位置编码向量和嵌入向量的维度是一样的,比如下边都是四个格子:
举个例子,当嵌入向量的长度为4的时候,位置编码长度也是4

一直说位置向量遵循某个模式,这个模式到底是什么。

参考论文:Convolutional Sequence to Sequence Learning (2017.05,-) 。

在下面的图中,每一行对应一个位置编码。所以第一行就是我们输入序列中第一个单词的位置编码,之后我们要把它加到词嵌入向量上。

看个可视化的图:
这里表示的是一个句子有20个词,词嵌入向量的长度为512。可以看到图像从中间一分为二,因为左半部分是由正弦函数生成的。右半部分由余弦函数生成。然后将它们二者拼接起来,形成了每个位置的位置编码。

你可以在get_timing_signal_1d() 中看到生成位置编码的代码。这不是位置编码的唯一方法。

但是需要注意注意一点,上图的可视化是官方Tensor2Tensor库中的实现方法,将sin和cos拼接起来。但是和论文原文写的不一样,论文原文的3.5节写了位置编码的公式,论文不是将两个函数concat起来,而是将sin和cos交替使用。论文中公式的写法可以看这个代码:transformer_positional_encoding_graph,其可视化结果如下:
这里表示的是一个句子有10个词,词嵌入向量的长度为64。

论文中的公式和Python实现见后。

PE(pos,2i)=sin(pos/100002i/dmodel)PE_{(pos,2i)} = \sin(pos / 10000^{2i/d_{\text{model}}}) PE(pos,2i+1)=cos(pos/100002i/dmodel)PE_{(pos,2i+1)} = \cos(pos / 10000^{2i/d_{\text{model}}})

其中 pospos 是位置,ii是维度。也就是说,位置编码的每个维度都对应一个正弦函数。波长形成一个从2π到10000⋅2π的几何级数。论文选择这个函数,是因为假设它可以让模型通过相对位置来轻松学习注意力,因为对于任何固定的偏移量 PEpos+kPE_{pos+k} 都可以表示为 PEposPE_{pos} 的线性函数。

class PositionalEncoding(nn.Module):

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

        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

    def forward(self, x):
        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
        return self.dropout(x)

全连接的前馈网络(Feed-Forward Networks)

除了注意力子层外,我们的编码器和解码器中的每一层都包含一个全连接的前馈网络,该网络被单独且相同地应用于每个位置。这包括两个线性变换,它们之间有ReLU激活函数。

FFN(x)=max(0,xW1+b1)W2+b2\mathrm{FFN}(x)=\max(0, xW_1 + b_1) W_2 + b_2

虽然FFN的网络架构在各个位置上都是相同的,但它们在每个位置使用的是不同的权重参数。这可能就是论文作者为了强调这个,加上PositionWise的原因。

另一种描述方法是,这是两个具有1内核大小的卷积。输入和输出的维度是dmodel=512d_{\text{model}}=512,而内部层的维度是 dff=2048d_{ff}=2048

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."

    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(self.w_1(x).relu()))

子层之间的连接(残差和层归一化)

子层之间的连接,原始论文和上述参考代码的实现有少许差异,后面分别说明。

原始论文

在继续往下讲之前,我们还需再提一下编码器层中的一个细节:每个编码器层中的每个子层(自注意力层、前馈神经网络)都有一个残差连接(图中的Add),之后是做了一个层归一化(layer-normalization)(图中的Normalize)。

将过程中的向量相加和layer-norm可视化如下所示:

参考论文:Layer Normalization (2016.07,多伦多大学等) ,代码实现如下。

class LayerNorm(nn.Module):

    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

当然在解码器子层中也是这样的。

我们现在画一个有两个编码器和解码器的Transformer,那就是下图这样的:

参考代码

关于 LayerNorm 的位置,示例代码实现使用的 Pre-Layer Normalization(Pre_LN),Transformers 原始论文用的是 Post-Layer Normalization (Post-LN)。使用 Pre_LN 的好处,可以参考论文 On Layer Normalization in the Transformer Architecture (2020.02,微软亚洲研究院等) 。

(a) Post-LN Transformer layer; (b) Pre-LN Transformer layer.

解码器(Decoder)

现在我们已经介绍了编码器的大部分概念,因为Encoder的Decoder差不多,我们基本上也知道了解码器是如何工作的。那让我们直接看看二者是如何协同工作的。

解码器首先处理输入序列,将最后一个编码器层的输出转换为一组注意向量K和V。注意:参考实现中为直接用,见EncoderDecoder.forward,DecoderLayer.forward。

每个解码器层将在“encoder-decoder attention”层中使用编码器传过来的K和V,这有助于解码器将注意力集中在输入序列中的适当位置:

完成编码阶段后,我们开始进行解码阶段。在解码阶段每一轮计算都只往外蹦一个输出,在本例中是输出一个翻译之后的英语单词。

输出步骤会一直重复,直到遇到句子结束符 表明transformer的解码器已完成输出。

每一步的输出都会在下一个时间步喂给给底部解码器,解码器会像编码器一样运算并输出结果(每次往外蹦一个词)。

跟编码器一样,在解码器中我们也为其添加位置编码,以指示每个单词的位置。

解码器中的自注意力层和编码器中的不太一样: 在解码器中,自注意力层只允许关注已输出位置的信息。实现方法是在自注意力层的softmax之前进行mask,将未输出位置的信息设为极小值。

“encoder-decoder attention”层的工作原理和前边的多头自注意力差不多,但是Q、K、V的来源不用,Q是从下层创建的(比如解码器的输入和下层decoder组件的输出),但是其K和V是来自编码器最后一个组件的输出结果。

最后的线性层和softmax层

Decoder输出的是一个浮点型向量(512维),如何把它变成一个词?

这就是最后一个线性层和softmax要做的事情。

线性层就是一个简单的全连接神经网络,它将解码器生成的向量映射到logits向量中。假设我们的模型词汇表是10000个英语单词,它们是从训练数据集中学习的。那logits向量维数也是10000,每一维对应一个单词的分数。

然后,softmax层将这些分数转化为概率(全部为正值,加起来等于1.0),选择其中概率最大的位置的词汇作为当前时间步的输出。

这张图从下往上看,假设具体上的那个向量是解码器的输出,然后将其转换为最终输出的单词。

from torch.nn.functional import log_softmax

class Generator(nn.Module):
    "Define standard linear + softmax generation step."

    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()

        # 输入为模型维度大小d_model(论文中为512),输出为词典大小
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return log_softmax(self.proj(x), dim=-1)

嵌入向量(Embeddings)

Transformer 中有两个嵌入向量层,用于将token id转换为512维的向量。

  1. 调用encode前,将源字典中的token id转换为向量。
  2. 调用decode前,将目标字典中的token id转换为向量。
class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

训练过程

现在我们已经了解了Transformer的整个前向传播的过程,那我们继续看一下训练过程。在训练期间,未经训练的模型会进行相同的前向传播过程。由于我们是在有标记的训练数据集上训练它,所以我们可以将其输出与实际的输出进行比较。

为了便于理解,我们假设预处理阶段得到的词汇表只包含六个单词(“a”, “am”, “i”, “thanks”, “student”, **“”**)。

这个词汇表是在预处理阶段创建的,在训练之前就已经有了。

一旦我们定义好了词汇表,我们就可以使用长度相同的向量(独热码,one-hot 向量)来表示词汇表中的每个单词。例如,我们可以用以下向量表示单词“am”:

接下来让我们讨论一下模型的损失函数,损失函数是我们在训练阶段优化模型的指标,通过损失函数,可以帮助我们获得一个准确的、我们想要的模型。

损失函数

假设我们正要训练我们的模型。

假设现在是训练阶段的第一步,我们用一个简单的例子(一个句子就一个词)来训练模型:把 “merci” 翻译成 “thanks”
这意味着,我们希望输出是表示“谢谢”的概率分布。但由于这个模型还没有经过训练,所以目前还不太可能实现。

由于模型的参数都是随机初始化的,未经训练的模型为每个单词生成任意的概率分布。我们可以将其与实际输出进行比较,然后使用反向传播调整模型的权重,使输出更接近我们所需要的值。

如何比较两种概率分布?在这个例子中我们只是将二者相减。实际应用中的损失函数有交叉熵(cross-entropy)(参考:交叉熵损失函数(Cross Entropy Loss):图示+公式+代码Visual Information Theory)和相对熵(也称为 Kullback–Leibler 散度)等。

上述只是最最简单的一个例子。现在我们来使用一个短句子(一个词的句子升级到三个词的句子了),比如输入 “je suis étudiant” 预期的翻译结果为: “i am a student”

所以我们希望模型不是一次输出一个词的概率分布了,能不能连续输出概率分布,最好满足下边要求:

对应上面的单词表,我们可以看出这里的one-hot向量是我们训练之后想要达到的目标。

在足够大的数据集上训练模型足够长的时间后,假设我们生成的概率分布如下所示: 训练之后最终得到的结果

模型一次产生一个输出,我们如何获得我们想要的输出呢?一般有如下两种方法。

  1. 贪心算法(greedy decoding):模型每次都选择分布概率最高的位置,输出其对应的单词。
  2. 集束搜索(beam search):保留概率最高前两个单词(例如,“I”和“a”),然后在下一步继续选择两个概率最高的值,以此类推,在这里我们把束搜索的宽度设置为2,当然你也可以设置其他的束搜索宽度。

整个模型

# 一个标准的编码器-解码器架构
class EncoderDecoder(nn.Module):

    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    def forward(self, src, tgt, src_mask, tgt_mask):
        "Take in and process masked src and target sequences."
        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)

    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

# 整个模型,包括编码器,解码器,输入和输出字典的Embeddings,生成器
def make_model(
    src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1
):
    "Helper: Construct a model from hyperparameters."
    c = copy.deepcopy
    attn = MultiHeadedAttention(h, d_model)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    position = PositionalEncoding(d_model, dropout)
    model = EncoderDecoder(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
        Generator(d_model, tgt_vocab),
    )

    # This was important from their code.
    # Initialize parameters with Glorot / fan_avg.
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model

示例程序,方便更好的理解上面的代码。

def subsequent_mask(size):
    "Mask out subsequent positions."
    attn_shape = (1, size, size)
    subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
        torch.uint8
    )
    return subsequent_mask == 0

def inference_test():
    #31:模型输入的词典大小,21:模型输出的词典大小;2:MHA包含2个头
    test_model = make_model(31, 21, 2)
    test_model.eval()

    # 1个示例句子,长度为11,每个数字对应输入词典中的序号
    src = torch.LongTensor([[12,11, 10,9,8,7,6,5,4,3,2]])
    src_mask = torch.ones(1, 1, 11)

    # 编码器的输出,矩阵大小为(1, 11, 512)
    memory = test_model.encode(src, src_mask)
    
    ys = torch.zeros(1, 1).type_as(src)

    # 逐个生成字典中的编号。
    for i in range(20):
        # out 长度从 1 逐渐增加到 20,示例形状[1, 18, 512]
        out = test_model.decode(
            memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
        )

        # 选中输出字典中的编号 
        prob = test_model.generator(out[:, -1])
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.data[0]
        
        ys = torch.cat(
            [ys, torch.empty(1, 1).type_as(src.data).fill_(next_word)], dim=1
        )

    print("Example Untrained Model Prediction:", ys)

# 执行测试代码    
inference_test()

发布于:2023-09-05 07:18:19 描述有误?我来纠错