字节终面问Transformer，就很离谱...

Transformer 是目前 NLP 甚至是整个深度学习领域不能不提到的框架，同时大部分 LLM 也是使用其进行训练生成模型，所以 Transformer 几乎是目前每一个机器人开发者或者人工智能开发者不能越过的一个框架。

接下来本文将从顶层往下去一步步掀开 Transformer 的面纱。

Transformer 概述

Transformer 模型来自论文 Attention Is All You Need。

链接：https://arxiv.org/abs/1706.03762

在论文中最初是为了提高机器翻译的效率，它使用了 Self-Attention 机制和 Position Encoding 去替代 RNN。

后来大家发现 Self-Attention 的效果很好，并且在其它的地方也可以使用 Transformer 模型。并引出后面的 BERT 和 GPT 系列。

大家一般看到的 Transformer 框架如下图所示：

Transformer 模型概览

首先把模型看成一个黑盒，如下图所示，对于机器翻译来说，它的输入是源语言(法语)的句子，输出是目标语言(英语)的句子。

把黑盒子稍微打开一点，Transformer(或者任何的 NMT 系统)可以分成 Encoder 和 Decoder 两个部分，如下图所示。

再展开一点，Encoder 由很多结构一样的 Encoder 堆叠而成，Decoder 也是一样。如下图所示。

每一个 Encoder 的输入是下一层 Encoder 输出，最底层 Encoder 的输入是原始的输入(法语句子)；Decoder 也是类似，但是最后一层 Encoder 的输出会输入给每一个 Decoder 层，这是 Attention 机制的要求。

每一层的 Encoder 都是相同的结构，它由一个 Self-Attention 层和一个前馈网络(全连接网络)组成，如下图所示。

每一层的 Decoder 也是相同的结构，它除了 Self-Attention 层和全连接层之外还多了一个 Attention 层，这个 Attention 层使得 Decoder 在解码时会考虑最后一层 Encoder 所有时刻的输出。它的结构如下图所示。

Transformer 流程串联

Transformer 的串流需要 tensor 的加入，输入的句子需要通过 Embedding 把它变成一个连续稠密的向量，如下图所示。

Embedding 之后的序列会输入 Encoder，首先经过 Self-Attention 层然后再经过全连接层。

我们在计算 𝑧𝑖 时需要依赖所有时刻的输入 𝑥1,…,𝑥𝑛，这是可以用矩阵运算一下子把所有的 𝑧𝑖 计算出来的。

而全连接网络的计算则完全是独立的，计算 i 时刻的输出只需要输入 𝑧𝑖 就足够了，因此很容易并行计算。

下图更加明确地表达了这一点。图中 Self-Attention 层是一个大的方框，表示它的输入是所有的 𝑥1,…,𝑥𝑛，输出是 𝑧1,…,𝑧𝑛。

而全连接层每个时刻是一个方框(但不同时刻的参数是共享的)，表示计算 𝑟𝑖 只需要 𝑧𝑖。此外，前一层的输出 𝑟1,…,𝑟𝑛 直接输入到下一层。

Self-Attention 介绍

比如我们要翻译如下句子”The animal didn’t cross the street because it was too tired”(这个动物无法穿越马路，因为它太累了)。

这里的 it 到底指代什么呢，是 animal 还是 street？要知道具体的指代，我们需要在理解 it 的时候同时关注所有的单词，重点是 animal、street 和 tired，然后根据知识(常识)我们知道只有 animal 才能 tired，而 street 是不能 tired 的。

Self-Attention 用 Encoder 在编码一个词的时候会考虑句子中所有其它的词，从而确定怎么编码当前词。如果把 tired 换成 narrow，那么 it 就指代的是 street 了。

下图是模型的最上一层 Encoder 的 Attention 可视化图。这是 tensor2tensor 这个工具输出的内容。

我们可以看到，在编码 it 的时候有一个 Attention Head(后面会讲到)注意到了 Animal，因此编码后的 it 有 Animal 的语义。

下面我们详细的介绍 Self-Attention 是怎么计算的，首先介绍向量的形式逐个时刻计算，这便于理解，接下来我们把它写出矩阵的形式一次计算所有时刻的结果。

对于输入的每一个向量(第一层是词的 Embedding，其它层是前一层的输出)，我们首先需要生成 3 个新的向量 Q、K 和 V，分别代表查询(Query)向量、Key 向量和 Value 向量。

Q 表示为了编码当前词，需要去注意(attend to)其它(其实也包括它自己)的词，我们需要有一个查询向量。

而 Key 向量可以认为是这个词的关键的用于被检索的信息，而 Value 向量是真正的内容。

具体的计算过程如下图所示。比如图中的输入是两个词”thinking”和”machines”，我们对它们进行 Embedding(这是第一层，如果是后面的层，直接输入就是向量了)，得到向量 𝑥1,𝑥2。

接着我们用 3 个矩阵分别对它们进行变换，得到向量 𝑞1,𝑘1,𝑣1 和 𝑞2,𝑘2,𝑣2。比如 𝑞1=𝑥1𝑊𝑄，图中 𝑥1 的 shape 是 1x4，𝑊𝑄 是 4x3，得到的 𝑞1 是 1x3。

其它的计算也是类似的，为了能够使得 Key 和 Query 可以内积，我们要求 𝑊𝐾 和 𝑊𝑄 的 shape 是一样的，但是并不要求 𝑊𝑉 和它们一定一样(虽然实际论文实现是一样的)。

每个时刻 t 都计算出 𝑄𝑡,𝐾𝑡,𝑉𝑡 之后，我们就可以来计算 Self-Attention 了。以第一个时刻为例，我们首先计算 𝑞1 和 𝑘1,𝑘2 的内积，得到 score。

过程如下图所示：

接下来使用 softmax 把得分变成概率，注意这里把得分除以 8(𝑑𝑘) 之后再计算的 softmax，根据论文的说法，这样计算梯度时会更加稳定(stable)。

计算过程如下图所示：

接下来用 softmax 得到的概率对所有时刻的 V 求加权平均，这样就可以认为得到的向量根据 Self-Attention 的概率综合考虑了所有时刻的输入信息。

计算过程如下图所示：

这里只是演示了计算第一个时刻的过程，计算其它时刻的过程是完全一样的。

softmax 示例代码：

Python
import numpy as np
   def softmax(x):
       """Compute softmax values for each sets of scores in x."""
       # e_x = np.exp(x)  
       e_x = np.exp(x )  
       return e_x / e_x.sum()

if __name__ == '__main__':  
    
    x = np.array([-3, 2, -1, 0])
    res = softmax(x )
    print(res)                        # [0.0056533  0.83902451 0.04177257 0.11354962]

特别注意，以上过程是可以并行计算的。

Multi-Head Attention

论文还提出了 Multi-Head Attention 的概念。其实很简单，前面定义的一组 Q、K 和 V 可以让一个词 attend to 相关的词，我们可以定义多组 Q、K 和 V，它们分别可以关注不同的上下文。

计算 Q、K 和 V 的过程还是一样，不过现在变换矩阵从一组(𝑊𝑄,𝑊𝐾,𝑊𝑉)变成了多组(𝑊𝑄0,𝑊𝐾0,𝑊𝑉0) ，(𝑊𝑄1,𝑊𝐾1,𝑊𝑉1)。

如下图所示：

对于输入矩阵(time_step, num_input)，每一组 Q、K 和 V 都可以得到一个输出矩阵 Z(time_step, num_features)。

如下图所示：

但是后面的全连接网络需要的输入是一个矩阵而不是多个矩阵，因此我们可以把多个 head 输出的 Z 按照第二个维度拼接起来，但是这样的特征有一些多，因此 Transformer 又用了一个线性变换(矩阵 𝑊𝑂)对它进行了压缩。

这个过程如下图所示：

上面的步骤涉及很多步骤和矩阵运算，我们用一张大图把整个过程表示出来，如下图所示。

我们已经学习了 Transformer 的 Self-Attention 机制，下面我们通过一个具体的例子来看看不同的 Attention Head 到底学习到了什么样的语义。

从上面两图的对比也能看出使用多个 Head 的好处——每个 Head(在数据的驱动下)学习到不同的语义。

位置编码(Positional Encoding)

我们的目的是用 Self-Attention 替代 RNN，RNN 能够记住过去的信息，这可以通过 Self-Attention“实时”的注意相关的任何词来实现等价(甚至更好)的效果。

RNN 还有一个特定就是能考虑词的顺序(位置)关系，一个句子即使词完全是相同的但是语义可能完全不同，比如”北京到上海的机票”与”上海到北京的机票”，它们的语义就有很大的差别。

我们上面的介绍的 Self-Attention 是不考虑词的顺序的，如果模型参数固定了，上面两个句子的北京都会被编码成相同的向量。

但是实际上我们可以期望这两个北京编码的结果不同，前者可能需要编码出发城市的语义，而后者需要包含目的城市的语义。

而 RNN 是可以(至少是可能)学到这一点的。当然 RNN 为了实现这一点的代价就是顺序处理，很难并行。

为了解决这个问题，我们需要引入位置编码，也就是 t 时刻的输入，除了 Embedding 之外(这是与位置无关的)，我们还引入一个向量，这个向量是与 t 有关的，我们把 Embedding 和位置编码向量加起来作为模型的输入。

这样的话如果两个词在不同的位置出现了，虽然它们的 Embedding 是相同的，但是由于位置编码不同，最终得到的向量也是不同的。

位置编码有很多方法，其中需要考虑的一个重要因素就是需要它编码的是相对位置的关系。

比如两个句子：”北京到上海的机票”和”你好，我们要一张北京到上海的机票”。

显然加入位置编码之后，两个北京的向量是不同的了，两个上海的向量也是不同的了，但是我们期望 Query(北京1)Key(上海1)却是等于Query(北京2)Key(上海2)的。具体的编码算法我们在代码部分再介绍。

位置编码加入后的模型如下图所示：

一个具体的位置编码的例子如下图所示：

残差和归一化

每个 Self-Attention 层都会加一个残差连接，然后是一个 LayerNorm 层，如下图所示。

下图展示了更多细节：输入 𝑥1,𝑥2 经 self-attention 层之后变成 𝑧1,𝑧2，然后和残差连接的输入 𝑥1,𝑥2 加起来，然后经过 LayerNorm 层输出给全连接层。

全连接层也是有一个残差连接和一个 LayerNorm 层，最后再输出给上一层。

Decoder 和 Encoder 是类似的，如下图所示，区别在于它多了一个 Encoder-Decoder Attention 层，这个层的输入除了来自 Self-Attention 之外还有 Encoder 最后一层的所有时刻的输出。

Encoder-Decoder Attention 层的 Query 来自前面一层，而 Key 和 Value 则来自 Encoder 的输出。

此外在解码器的编码器-解码器注意力层中，掩码的使用非常关键，以确保解码器在生成每个目标词时只能使用到源语言句子的信息和它之前已经生成的目标词的信息。

Pytorch 实现 Transformer

代码如下：

Python
import torch
import torch.nn as nn
import math

# 位置编码模块
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()       
        pe = torch.zeros(max_len, d_model)       
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)       
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-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(0), :
        return x

# Transformer模型
class TransformerModel(nn.Module):   
    def __init__(self, ntoken, d_model, nhead, d_hid, nlayers, dropout=0.5):    
        super(TransformerModel, self).__init__()
        self.model_type = 'Transformer'
        self.pos_encoder = PositionalEncoding(d_model)
        self.encoder = nn.Embedding(ntoken, d_model)
        self.transformer = nn.Transformer(d_model, nhead, d_hid, nlayers, dropout)
        self.decoder = nn.Linear(d_model, ntoken)
        self.init_weights()
        self.dropout = nn.Dropout(dropout) 
        
    def generate_square_subsequent_mask(self, sz):  
       `# 生成后续掩码，用于防止位置信息泄露        
       mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)        
       mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))        
       return mask  
       
    def init_weights(self):  
        # 初始化权重
        initrange = 0.1
        self.encoder.weight.data.uniform_(-initrange, initrange)`        
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-initrange, initrange)
       
    def forward(self, src, src_mask)
        # 前向传播      
        src = self.encoder(src) * math.sqrt(self.d_model)        
        src = self.pos_encoder(src)        
        output = self.transformer(src, src, src_key_padding_mask=src_mask)       
        output = self.decoder(output)        
        return output

# 示例使用
ntokens = 1000  # 词汇表大小
d_model = 512  # 嵌入维度
nhead = 8  # 多头注意力中的头数
d_hid = 2048  # 前馈网络模型的维度
nlayers = 6  # 层数
dropout = 0.2  # dropout比率

model = TransformerModel(ntokens, d_model, nhead, d_hid, nlayers, dropout)

# 示例输入
src = torch.randint(0, ntokens, (10, 32))  # (序列长度, 批量大小)
src_mask = model.generate_square_subsequent_mask(10)  # 创建掩码

output = model(src, src_mask)
print(output)

推理过程

在 Transformer 模型的机器翻译任务中，解码器生成第一个翻译后的词（通常称为第一个目标词）的过程如下：

1、起始符号： 在解码器的输入序列的开始位置，通常会添加一个特殊的起始符号，如（Start Of Sentence）。这个符号告诉模型翻译过程的开始。

2、初始化隐藏状态： 解码器的隐藏状态通常初始化为零向量或从编码器的最后一层的输出中获得。这个隐藏状态在生成序列的每一步中都会更新。

3、第一次迭代： 在第一次迭代中，解码器的输入只包含起始符号。

解码器通过以下步骤生成第一个词：