导言

在语言模型取得ELMo和ULMFit等突破之后，转换器席卷了自然语言处理（NLP）领域。它是诸如GPT-3和DALL-E之类的流行语言模型的基础。同样，HuggingFace Transformers库之类的工具使机器学习工程师可以轻松地解决各种NLP任务，并促进了许多NLP的后续突破。

在此报告中，我们将深入研究论文“Attention Is All You Need”（2017） 中所述的转换器架构，并使用PyTorch对其进行编码。

转换器是一个序列到序列（seq2seq）模型。这意味着它会适用于数据中存在某些排序并且输出序列本身的这种问题。示例应用是机器翻译，抽象总结和语音识别。最近，Vision Transformers (ViT) 甚至改进了计算机视觉的最新技术。

在下面，您可以看到完整的转换器架构的可视化。我们将解释每个组件的作用、存在的原因以及所有组件的组合方式。现在，请注意这里有一个编码器（在左侧）和一个解码器（在右侧），它们内部都具有多个神经网络层。

示例代码是从Harvard NLP Group的The Annotated Transformer 和有关转换器的PyTorch文档中重构的。

分词(Tokenization)

首先，我们需要一种数字表示文本的方法，以便对其进行计算。分词是将文本字符串解析为压缩符号序列的过程。此过程会产生一个整数向量，其中每个整数代表文本的一部分。这篇转换器论文使用Byte-Pair Encoding (BPE)作为分词方法。 BPE是一种压缩形式，其中最常见的连续字节（即字符）被压缩为单个字节（即整数）。

最近的研究表明BPE并非最佳选择，而像BERT这样的最新语言模型则使用WordPiece分词器。 WordPiece似乎更易于解码也和更直观，因为它经常将整个单词标记为一个token。相比之下，BPE的分词则经常不完整。这会导致奇怪的token和单个字符（即字母表中的字母）被标记为未知token。因此，WordPiece是我们在示例中使用的分词器。我们可以在文本语料库上训练自己的分词器，但实际上，从业人员几乎总是使用预先训练好的分词器。 HuggingFace转换器库让预训练分词器的使用变得容易。

>>> from transformers import BertTokenizer
>>> tok = BertTokenizer.from_pretrained("bert-base-uncased")
Downloading: 100%|█████████████████████████| 1.04M/1.04M [00:00<00:00, 1.55MB/s]
Downloading: 100%|████████████████████████████| 456k/456k [00:00<00:00, 848kB/s]

>>> tok("Hello, how are you doing?")['inputs_ids']
{'input_ids': [101, 7592, 2129, 2024, 2017, 2725, 1029, 102]}

>>> tok("The Frenchman spoke in the [MASK] language and ate 🥖")['input_ids']
{'input_ids': [101, 1996, 26529, 3764, 1999, 1996, 103, 2653, 1998, 8823, 100, 102]}

>>> tok("[CLS] [SEP] [MASK] [UNK]")['input_ids']
{'input_ids': [101, 101, 102, 103, 100, 102]}

请注意，分词器会自动包含一个token用于编码开始($[CLS]=101$)和编码结束（即separation）（($[SEP]=102$)。 其他特殊token包括masking ($[MASK]=103$)和未知符号($[UNK]=100$)。

词嵌入(Embeddings)

为了学习文本的正确表示，序列中每个单独的token都通过嵌入转换为向量。可以将其视为一种神经网络层，因为嵌入的权重是与转换器模型的其余部分一起学习的。 它包含词汇表中每个单词的向量，并且这些权重是用正态分布N $\mathcal{N}(0,1)$.初始化的。 请注意，在初始化模型时($E\in {\mathbb{R}}^{vocab\times d_{model}}$)，它要求我们指定词汇表大小 ($vocab$) 和模型维度($d_{model}=512$) 。最后，权重乘以$\sqrt{d_{model}}$ 作为标准化步骤。

import torch
from torch import nn
class Embed(nn.Module):
    def __init__(self, vocab: int, d_model: int = 512):
        super(Embed, self).__init__()
        self.d_model = d_model
        self.vocab = vocab
        self.emb = nn.Embedding(self.vocab, self.d_model)
        self.scaling = math.sqrt(self.d_model)

    def forward(self, x):
        return self.emb(x) * self.scaling

位置编码（Positional Encoding）

与递归和卷积网络相反，该模型本身不具有序列中嵌入token的相对位置的信息。 因此，我们必须为编码器和解码器的输入嵌入中加入此信息——添加一个编码。 可以以许多不同的方式添加此信息，这些信息可以是静态的也可以是训练学习的。 转换器对每个位置($pos$)使用正弦和余弦变换。 正弦用于偶数维($2_i$)，余弦用于奇数维($2_{i+1}$)。

$P{E}_{pos,2i}=sin(\frac{pos}{10000^{2i/d_{model}}})$

$P{E}_{pos,2i+1}=cos(\frac{pos}{10000^{2i/d_{model}}})$

在代码中，在对数空间中计算位置编码以避免数值溢出。

import torch
from torch import nn
from torch.autograd import Variable

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int = 512, dropout: float = .1, max_len: int = 5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        
        # Compute the positional encodings 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) * -(torch.log(torch.Tensor([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 + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        return self.dropout(x)

多头注意力(Multi-Head Attention)

在转换器出现之前，人工智能研究中从序列中学习习惯使用卷积（(WaveNet，ByteNet）或递归（(RNN, LSTM，LSTM）。在转换器之前，注意力(Attention)已经可以实现一些自然语言处理方面的突破（Luong等人，2015），但是当时还不清楚其实不需要卷积或递归来构建有效的模型。因此，提出“Attention is All You Need”是一个大胆的声明。

注意力层可以学习query($Q$)和一组key($K$) value($V$) 之间的映射。这些名称的含义可能会让人迷惑，因为它们取决于特定的NLP应用。在文本生成中，query是输入的词嵌入。value和key可以视为目标。通常，value和key是相同的。

例如，当您在Youtube上搜索视频时，您将在搜索栏中键入一个短语（即query $Q$）。搜索引擎将使用它来映射到Youtube视频的标题、描述等（即key $K$）。使用此映射，它将向您建议最相关的视频（即value $V$）。 （示例来源）

提高NLP中注意力表现的一个创新——被作者称为"Scaled dot-product attention"。它与multiplicative attention相同，但$Q$和$K$的映射被key的维度$d_k$缩放。这使得multiplicative attention在较大的维度上表现更好。结果会被通过softmax激活方程（($softmax(x_i)=\frac{\exp (x_i)}{{\sum }_{j}exp(x_j)}$) a并乘以$V$。

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

import torch
from torch import nn
class Attention:
    def __init__(self, dropout: float = 0.):
        super(Attention, self).__init__()
        self.dropout = nn.Dropout(dropout)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, query, key, value, mask=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 = self.dropout(self.softmax(scores))
        return torch.matmul(p_attn, value)
    
    def __call__(self, query, key, value, mask=None):
        return self.forward(query, key, value, mask)

在解码器中，通过用很大的负数（$-1e9$或有时甚至是-$-inf$）填充某些位置来掩盖attention子层。 这是为了防止模型关注后续位置来作弊。 这样可以确保模型在尝试预测下一个token时只能关注先前位置的单词。

注意力机制本身已经非常强大，可以在现代硬件（例如针对矩阵乘法进行了优化的GPU和TPU）上有效地进行计算。 但是，单个attention层仅允许一个表征。 因此，在转换器中使用了多个attention。 这使模型可以学习多种模式和表征。 本文使用h = 8个串联的attention层。 最终公式变为：

$MultiHead(Q,K,V)=Concat(hea{d}_1,...,hea{d}_n)W^O$

where $hea{d}_i=Attention(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$

The projection weights ($W^O\in {\mathbb{R}}^{h{d}_v\times d_{model}},W_i^Q\in {\mathbb{R}}^{d_{model}\times d_k},W_i^K\in {\mathbb{R}}^{d_{model}\times d_k},W_i^V\in {\mathbb{R}}^{d_{model}\times d_v}$) are outputs of a fully-connected (Linear) layer. The authors of the transformer paper use $d_k=d_v=\frac{d_{model}}{h}=64$. 预测的权重[($W^O\in {\mathbb{R}}^{h{d}_v\times d_{model}},W_i^Q\in {\mathbb{R}}^{d_{model}\times d_k},W_i^K\in {\mathbb{R}}^{d_{model}\times d_k},W_i^V\in {\mathbb{R}}^{d_{model}\times d_v}$]是完全连接的（线性）层的输出。 转换器论文的作者使用[function]

from torch import nn
from copy import deepcopy
class MultiHeadAttention(nn.Module):
    def __init__(self, h: int = 8, d_model: int = 512, dropout: float = 0.1):
        super(MultiHeadAttention, self).__init__()
        self.d_k = d_model // h
        self.h = h
        self.attn = Attention(dropout)
        self.lindim = (d_model, d_model)
        self.linears = nn.ModuleList([deepcopy(nn.Linear(*self.lindim)) for _ in range(4)])
        self.final_linear = nn.Linear(*self.lindim, bias=False)
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, query, key, value, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(1)
        
        query, key, value = [l(x).view(query.size(0), -1, self.h, self.d_k).transpose(1, 2) \
                             for l, x in zip(self.linears, (query, key, value))]
        nbatches = query.size(0)
        x = self.attn(query, key, value, mask=mask)
        
        # Concatenate and multiply by W^O
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
        return self.final_linear(x)

• *技术说明：.conpose是在.transpose之后添加的，因为.transpose与原始张量共享其基础内存存储。 之后调用.view需要一个连续的张量（文档）。 .view方法允许有效地重塑、切片和按元素操作（文档）。

  由于每个attention头的维度都除以$h$，因此总计算类似于使用一个具有完整维度的attention头($d_{model}$)。 但是，通过这种方法，计算可以沿每个头并行进行，这实现了现代硬件的大规模加速。 这是创新之一——不使用卷积或递归但训练有效的语言模型。

残差和网络层标准化

AI研究社区发现，诸如残差连接residual connections和（批）标准化 (batch) normalization之类的概念可以提高性能，减少训练时间并允许训练更深的网络。因此，在每个attention层和每个前馈层(feed forward layer)之后，转换器都具有残差连接和标准化功能。另外，为了更好的一般化，在每层中都添加了dropout。

标准化

基于现代深度学习的计算机视觉模型通常具有批标准化的功能。但是，这种类型的标准化依赖于大批量，并且无法自己递归。传统的转换器架构改为具有层标准化layer normalization。即使批量大小较小($batchsize<8$)，层标准化也是稳定的。

为了计算层标准化，我们首先分别为minibatch中的每个样本计算平均值${\mu }_i$和标准差 ${\sigma }_i$。

${\mu }_i=\frac{1}{K}{\sum }_{k=1}^k{x}_{i,k}$

${\sigma }_i=\sqrt{\frac{1}{K}{\sum }_{k=1}^k(x_{i,k}-{\mu }_i{)}^2}$

然后，将标准化步骤定义为：

$L{N}_{\gamma ,\beta }(x_i)\equiv \gamma \frac{x-{\mu }_i}{{\sigma }_i+ϵ}+\beta$

其中$\gamma$ 和 $\beta$是可学习的参数。在标准偏差$ϵ$为$0$的情况下，为数字稳定性添加了一个小数${\sigma }_i$。

from torch import nn
class LayerNorm(nn.Module):
    def __init__(self, features: int, eps: float = 1e-6):
        super(LayerNorm, self).__init__()
        self.gamma = nn.Parameter(torch.ones(features))
        self.beta = 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.gamma * (x - mean) / (std + self.eps) + self.beta

残差

残差连接意味着您将网络（即子层）中上一层的输出添加到当前层的输出中。 这允许非常深的网络，因为该网络实际上可以“跳过”某些层。

每层的最终输出将是 $ResidualConnection(x)=x+Dropout(SubLayer(LayerNorm(x)))$

from torch import nn
class ResidualConnection(nn.Module):
    def __init__(self, size: int = 512, dropout: float = .1):
        super(ResidualConnection,  self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))

前馈 (Feed Forward)

在每个attention层的顶部都添加了前馈网络。 它由两个完全连接的层组成，这些层具有$ReLU$激活函数 ($ReLU(x)=max(0,x)$) 和内层的dropout。 转换器论文中使用的标准维度对于输入层是$d_{model}=512$，对于内层来说是$d_{ff}=2048$。

完整的计算变成 $FeedForward(x)=W_{2}max(0,x{W}_1+B_1)+B_2$

请注意，默认情况下，PyTorch Linear 已包含偏差（B1和B2）。

from torch import nn
class FeedForward(nn.Module):
    def __init__(self, d_model: int = 512, d_ff: int = 2048, dropout: float = .1):
        super(FeedForward, self).__init__()
        self.l1 = nn.Linear(d_model, d_ff)
        self.l2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.l2(self.dropout(self.relu(self.l1(x))))

编码器-解码器

编码

现在，我们具有构建模型编码器和解码器的所有组件。 单个编码器层由一个多头attention层和一个前馈网络组成。 如前所述，我们还包括残差连接和层标准化。

$Encoding(x,mask)=FeedForward(MultiHeadAttention(x))$

from torch import nn
from copy import deepcopy
class EncoderLayer(nn.Module):
    def __init__(self, size: int, self_attn: MultiHeadAttention, feed_forward: FeedForward, dropout: float = .1):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sub1 = ResidualConnection(size, dropout)
        self.sub2 = ResidualConnection(size, dropout)
        self.size = size

    def forward(self, x, mask):
        x = self.sub1(x, lambda x: self.self_attn(x, x, x, mask))
        return self.sub2(x, self.feed_forward)

论文的最终转换器编码器由6个相同的编码器层组成，然后进行层标准化。

class Encoder(nn.Module):
    def __init__(self, layer, n: int = 6):
        super(Encoder, self).__init__()
        self.layers = nn.ModuleList([deepcopy(layer) for _ in range(n)])
        self.norm = LayerNorm(layer.size)
        
    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

解码

解码层是一个masked multi-head attention layer（多头注意力层），其后是包括memory的multi-head attention layer（多头注意力层）。memory是编码器的输出。 最后，它会通过前馈网络。 同样，所有这些组件都包括残差连接和层标准化。

$Decoding(x,memory,mask1,mask2)=FeedForward(MultiHeadAttention(MultiHeadAttention(x,mask1),memory,mask2))$

from torch import nn
from copy import deepcopy
class DecoderLayer(nn.Module):
    def __init__(self, size: int, self_attn: MultiHeadAttention, src_attn: MultiHeadAttention, 
                 feed_forward: FeedForward, dropout: float = .1):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sub1 = ResidualConnection(size, dropout)
        self.sub2 = ResidualConnection(size, dropout)
        self.sub3 = ResidualConnection(size, dropout)
 
    def forward(self, x, memory, src_mask, tgt_mask):
        x = self.sub1(x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sub2(x, lambda x: self.src_attn(x, memory, memory, src_mask))
        return self.sub3(x, self.feed_forward)

与最终编码器一样，论文中的解码器也具有6个相同的层，然后进行层标准化。

class Decoder(nn.Module):
    def __init__(self, layer: DecoderLayer, n: int = 6):
        super(Decoder, self).__init__()
        self.layers = nn.ModuleList([deepcopy(layer) for _ in range(n)])
        self.norm = LayerNorm(layer.size)
        
    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)

通过编码器和解码器的这种更高级别的表示，我们可以轻松地制定最终的编码器-解码器块。

from torch import nn
class EncoderDecoder(nn.Module):
    def __init__(self, encoder: Encoder, decoder: Decoder, 
                 src_embed: Embed, tgt_embed: Embed, final_layer: Output):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.final_layer = final_layer
        
    def forward(self, src, tgt, src_mask, tgt_mask):
        return self.final_layer(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)

最终输出

最后，必须将解码器的矢量输出转换为最终输出。对于像语言翻译这样的序列到序列问题，这是每个位置在整个词汇量上的概率分布。一个完全连接的层将解码器的输出转换为logits矩阵，该矩阵具有目标词汇表的维度。这些数字通过softmax激活函数转换为词汇表上的概率分布。在代码中，我们使用$LogSoftmax(x_i)=log(\frac{\exp (x_i)}{{\sum }_{j}exp(x_j)})$，因为它更快并且在数值上更稳定。

例如，假设翻译后的句子有$20$个token，总词汇量为$30000$个token。输出结果将是矩阵$M\in {\mathbb{R}}^{20\times 30000}$。然后，我们可以在最后一个维度上使用$argmax$以获得输出token $T\in {\mathbb{N}}^{20}$ 的向量，该向量可以通过分词器解码为文本字符串。

$Output(x)=LogSoftmax(max(0,x{W}_1+B_1))$

from torch import nn
class Output(nn.Module):
    def __init__(self, input_dim: int, output_dim: int):
        super(Output, self).__init__()
        self.l1 = nn.Linear(input_dim, output_dim)
        self.log_softmax = nn.LogSoftmax(dim=-1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        logits = self.l1(x)
        return self.log_softmax(logits)

模型初始化

我们以与论文相同的维度构建转换器模型。初始化策略是Xavier / Glorot初始化，从$[-\frac{1}{\sqrt{n}},\frac{1}{\sqrt{n}}]$范围内的均匀分布中选取。所有偏差均初始化为$0$。

$Xavier(W)\sim U[-\frac{1}{\sqrt{n}},\frac{1}{\sqrt{n}}],B=0$

from torch import nn
def make_model(input_vocab: int, output_vocab: int, d_model: int = 512):
    encoder = Encoder(EncoderLayer(d_model, MultiHeadAttention(), FeedForward()))
    decoder = Decoder(DecoderLayer(d_model, MultiHeadAttention(), MultiHeadAttention(), FeedForward()))
    input_embed= nn.Sequential(Embed(vocab=input_vocab), PositionalEncoding())
    output_embed = nn.Sequential(Embed(vocab=output_vocab), PositionalEncoding())
    output = Output(input_dim=d_model, output_dim=output_vocab)
    model = EncoderDecoder(encoder, decoder, input_embed, output_embed, output)
    
    # Initialize parameters with Xavier uniform 
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model

该函数返回一个PyTorch模型，该模型可以针对序列问题进行训练。下面是一个示例，说明如何将其与分词的输入和输出一起使用。假设我们的输入和输出词汇只有$10$个单词。

# Tokenized symbols for source and target.
>>> src = torch.tensor([[1, 2, 3, 4, 5]])
>>> src_mask = torch.tensor([[1, 1, 1, 1, 1]])
>>> tgt = torch.tensor([[6, 7, 8, 0, 0]])
>>> tgt_mask = torch.tensor([[1, 1, 1, 0, 0]])

# Create PyTorch model
>>> model = make_model(input_vocab=10, output_vocab=10)
# Do inference and take tokens with highest probability through argmax along the vocabulary axis (-1)
>>> result = model(src, tgt, src_mask, tgt_mask)
>>> result.argmax(dim=-1)
tensor([[6, 6, 4, 3, 6]])

输出与目标相差很远，因为此时模型具有统一初始化的权重。从零开始训练这些转换器模型需要大量的计算。为了训练基本模型，原始论文的作者在8个NVIDIA P100 GPU上进行了12小时的训练。他们的大型模型花了3.5天的时间在8个GPU上进行训练！我建议使用预训练的转换器模型，并针对您的应用对其进行微调。 HuggingFace Transformers库已经具有许多用于微调的预训练模型。

如果您确实想了解有关从头开始编写训练程序的更多信息，我建议您查看Annotated Transformer的训练部分。

就这样！希望您喜欢这次深入转换器架构的学习！

如果您有任何疑问或反馈，请在下面发表评论。您也可以通过Twitter @carlolepelaars与我联系。

数据来源和更多资源

推荐论文

Attention Is All You Need (2017)

BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (2018)

Reformer: The Efficient Transformer (2020)

Linformer: Self-Attention with Linear Complexity (2020)

Longformer: The Long-Document Transformer (2020)

Language Models are Few-Shot Learners (GPT-3 paper)(2020)

网上资源

The Annotated transformer with PyTorch Code

The Illustrated transformer

The Narrated transformer video

Łukasz Kaiser's masterclass on transformers