迟到的transformer encoder代码详解

前言

与传统序列模型不同，transformer的创新点在于能够捕捉语义全局信息(同时通过position embedding考虑到了序列之间的位置关系)、能够并行化计算…
想通过本文的代码层面的记录，让我和大家一眼就可以知道(或者记起)transformer模型的架构以及实现方法。但背后究竟是什么原理，本文没有深究。

从“TransformerEncoder” 类说起

这个类实现了transformer的encoder的所有功能:

• word_embedding(1) + position_embedding(2)
• 计算attn_pad_mask,用于后续的softmax(3)
• 循环"n_layers"次单层"encoder layer"，得到最后的输出(4)
• 最后套一个线性层、激活层，用作任务的输出(5)(不一定)

首先:该类的输入名叫inputs，大小为 batch_size, seq_len

(1) word_embedding

原码中的word_embedding层是随机初始化的，其中的d_model在原文中是512

self.embedding = nn.Embedding(vocab_size,d_model)

但我也可以使用预训练的词向量模型，并决定其是否参与反向传播

self.embedding.weight = nn.Embedding.from_pretrained(embedding_weight)
self.embedding.weight.requires_grad = False if fix_embedding else True # 设定emb是否随训练更新

一般，模型的输入的大小为: batch_size, seq_len，经过word_embedding，就变成了 batch_size, seq_len. d_model

word_embedding = self.embedding(inputs) # batch_size, seq_len. d_model

除了word_embedding还不够，我们还需要position_embedding来得到序列关系

(2) positional_embedding(positional_encoding)

这部分其实只做了一件事:建立了一个叫做"sinusoid_table"的表，其大小为: seq_len+1,d_model。为什么要seq_len+1？—— 因为第0个位置向量是用来表示"pad"的(如果你的pad_id = 0的话)。具体看看它是怎么创建的:

def get_sinusoid_table(self, seq_len, d_model):
    def get_angle(pos, i, d_model):
        return pos / np.power(10000, (2 * (i // 2)) / d_model)
    sinusoid_table = np.zeros((seq_len, d_model))
    for pos in range(seq_len):
        for i in range(d_model):
            if i % 2 == 0:
                sinusoid_table[pos, i] = np.sin(get_angle(pos, i, d_model))
            else:
                sinusoid_table[pos, i] = np.cos(get_angle(pos, i, d_model))
    return torch.FloatTensor(sinusoid_table)

这与原文中的公式👇是一致的:

有了这个表，当然是要跟word_embedding一样，每一个单词都要去表中索引到自己的向量表示，具体实现和思路如下:

positions = torch.arange(inputs.size(1), device=inputs.device, dtype=inputs.dtype).repeat(inputs.size(0), 1) + 1
# 为什么要加1呢? 1. sinusoid_table中也加了1，所以不会报错  2. 第0个位置是留给pad的
position_pad_mask = inputs.eq(self.pad_id)
# inputs中,值为0的(代表是pad)在position_pad_mask中都为1,非0的(非pad的)都为0
positions.masked_fill_(position_pad_mask, 0)
# positions中,非pad的位置不会被填充0,它们依然是原来在句子中的位置,而pad的位置都改成了0,之后会索引到sinusoid_table中的第一行

使用positions就可以去索引得到 batch_size, seq_len, d_model的位置表示矩阵:

position_embeding = self.pos_embedding(positions)

然后将word_embedding 与 position_embedding求和，得到的词向量表示记为 input_embedding(batchsize, seq_len,d_model)

(3) 计算attn_pad_mask

attn_pad_mask是什么?
——在multiheadattention部分, $Q\ast K^T/\sqrt{(}{d}_{model})$之后会得到 batch_size, seq_len, seq_len大小的权重矩阵(记作W)，其中的$W_{i,j}$表示第i个单词与第j个单词之间的关联度,而我们需要所有单词与第i个单词的关联度向量(记做$W_i$)。
用来干嘛?———在self attention部分, 第i个单词的向量就是表示为: $\sum W_{i,j}\ast Ve{c}_j(j=0,1,....)$其中的$Ve{c}_j$是第j个单词的向量表示(这就是注意力机制)。
而attn_pad_mask的作用就是标记所有pad的位置，以便让$W_{i,pad}=0(i=0,1,....)$总而消除padding单词对其它单词的影响。
说了这么多，怎么得到呢?

def get_attention_padding_mask(self, q, k, pad_id):
    attn_pad_mask = k.eq(pad_id).unsqueeze(1).repeat(1, q.size(1), 1)
    # |attn_pad_mask| : (batch_size, q_len, k_len)
    return attn_pad_mask
attn_pad_mask = self.get_attention_padding_mask(inputs, inputs, self.pad_id)

(这里的q、k都是inputs

(4) 循环"n_layers"次“EncoderLayer”类

这里引出EncoderLayer层，它就是单层完整的encoder，而transfomer encoder就是encoderLayer的循环。
该层的输入就是:上面word_embedding+position_embedding得到的input_embedding，以及attn_pad_mask。

"EncoderLayer"类——transformer encoder部分的核心

先放模型图:

其中，底层的 word_embedding和positional embedding之前已经完成了，接下来的部分从下到上(图)就分为四个部分:

• Muti-Head Attention(多头注意力机制)(1)(它是一个额外定义的类)
• MHA的输出+MHA的输入，去做一个Layer normalization(2)
• FeedForwardNetwork(前馈神经网络)(3)(它是一个额外定义的类)
• FFN的输出+FFN的输入，去做一个Layer normalization(4)

再回顾一下:该类的输入是word_embedding+position_embedding得到的input_embedding(batch_size, seq_len, d_model)，以及attn_pad_mask(batch_size, seq_len, seq_len)。

(1) Multi-Head Attention(MHA)(是一个额外定义的类)

不想让文章显得太复杂，所以就在这里完整介绍多头注意力机制。先说transformer里的注意力机制，再谈多头。
关于注意力机制，在上文的attn_pad_mask中已经提过。其核心思想就是:用其它所有单词的词向量表示，通过加权求和，来表示其中的一个单词(记作单词i)。而其中的"权"，说白了就是单词i的词向量表示分别与其它所有单词的词向量表示所做的"点积"——Dot-product(点积的实现又是另外一个类，下面的代码是从不同类之间拼接而成)。
具体怎么实现?
transformer选择的做法是:

1. 通过线性层，得到Q、K、V

把该类的第一个输入(word_embedding+position_embedding得到的input_embedding(batch_size, seq_len, d_model))，复制成三份，分别通过三个权重不共享、但规格相同的线性层，得到三个规格相同的矩阵Q、K、V。其中，三个权重的大小都为:(d_model,d_model)，所以得到的三个矩阵大小都保持不变、也都相同，但QKV各自发挥的作用从此就要分道扬镳了，代码如下

q_heads = self.WQ(inputs).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

k_heads = self.WK(inputs).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

v_heads = self.WV(inputs).view(batch_size, -1, self.n_heads, self.d_v).transpose(1, 2)

虽然上面还有self.n_heads这个东西，而且变量的命名也不是:Q\K\V，但完全可以先忽视它们，就当代码是这样写的，因为多头的思想在理解了注意力机制之后，就很好理解。:

Q = self.WQ(inputs)  # 依然是batch_size, seq_len, d_model
 
K = self.WK(inputs)  # 依然是batch_size, seq_len, d_model

V = self.WV(inputs)  # 依然是batch_size, seq_len, d_model

2. 计算$Q\ast K^T/\sqrt{(}{d}_{model})$

已知Q为batch_size, seq_len, d_model, $K^T$为batch_size,d_model,seq_len(transpose(-1,-2))。它们两个做相乘的话，得到的矩阵(记作W)规格为:batch_size, seq_len, seq_len。这个矩阵有什么含义?——就拿第i行来说(忽略batch_size)，$W_i$是一个seq_len长度的向量，里面的每一个$W_{ij}$就是第i个单词与第j个单词，两者词向量表示的点积(表示两个单词之间的联系)，而点积的实现，是通过K矩阵整体的转置来实现的。而$\sqrt{(}{d}_{model})$的解释，在原文中也给出了:

代码如下:

attn_score = torch.matmul(q_heads, k_heads.transpose(-1, -2)) / np.sqrt(self.d_k)  # d_k = d_model//n_heads

依然不考虑多头，改代码为:

attn_score = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(self.d_model)

3. 对上面计算得到的W做softmax归一化，得到真正的权重。

更直接点，就是把所有其它单词对于单词i的权重，做归一化。但其中的一部分单词是padding过来的，不具备实际意义。所以要用到之前记录好的attn_pad_mask(batch_size, seq_len, seq_len)，该矩阵中，所有padding的位置都为1，非padding的为0。要做的就是:在softmax之前，用masked_fill_函数以及这个attn_pad_mask，把上面的W中所有padding位置上的值改成负无穷，这样在softmax之后，那个位置的权重就是0，这样padding的单词就不会影响到其它非padding的单词。

attn_score.masked_fill_(attn_mask, -1e9)

然后就可以softmax了

attn_weights = nn.Softmax(dim=-1)(attn_score)

atten_weights的规格为: batch_size, seq_len, seq_len

4. 计算$attenweights\ast V$

V与Q、K一样，都是batch_size, seq_len, d_model大小的。记相乘得到的矩阵为output,它的大小为:batch_size, seq_len, d_model。忽略batch_size，就拿现在output中的第i个单词而言，它的表示依旧是一个 d_model长度的向量(记作$W_i$)，但这个向量的d_model个维度上，第k个维度的值= ${\sum }_{j=1}^{seqlen}attnweight{s}_{i,j}\ast V_{j,k}$, 也就是用其它向量对应维度的值做了加权求和所得。这便是attention的思想。

output = torch.matmul(attn_weights, V) # batch_size, seq_len, d_model

再来说说多头:就是在最最开始，就把Q、K、V在d_model这个维度上，再分成n_heads个头，得到的q_heads，k_heads, v_heads的大小都是 batch_size, seq_len, n_heads, d_model//n_heads， 再通过transpose(1,2)，得到batch_size, n_heads, seq_len, d_model//n_heads。如果我们省略前两维，剩下的就是 seq_len, d_model//n_heads这两个维度，我们依然可以做上面的注意力机制，只是原来的 d_model 变成了 d_model//n_heads。

q_heads = self.WQ(inputs).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
k_heads = self.WK(inputs).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
v_heads = self.WV(inputs).view(batch_size, -1, self.n_heads, self.d_v).transpose(1, 2) # d_v = d_k
attn_mask = attn_mask.unsqueeze(1).repeat(1, self.n_heads, 1, 1)	
attn_score = torch.matmul(q_heads, k_heads.transpose(-1, -2)) / np.sqrt(self.d_k)
attn_score.masked_fill_(attn_mask, -1e9)
attn_weights = nn.Softmax(dim=-1)(attn_score)
output = torch.matmul(attn_weights, v)  # batch_size, n_heads, seq_len, d_model//n_heads
output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.n_heads * self.d_v) 

最后一句代码把output拼接回: batch_size, seq_len, d_model,与不做多头的输出规格一样。
为什么用多头？——斯认为多头的作用有二(待考证)

• 增加计算的并行性
• 不同的头可以学习到不同的信息

(2) MHA的输出(output)+MHA(input_embedding)的输入，去做一个Layer normalization

self.layernorm1 = nn.LayerNorm(d_model, eps=1e-6)
attn_outputs = self.layernorm1(inputs + attn_outputs)  # batch_size, seq_len, d_model

关于layer normalization我一直没有很理解。

(3) FeedForwardNetwork(FFN)(是一个额外定义的类)

这个简单，FFN的公式如下:

说白了就是:

• 线性层1
• relu激活层
• 线性层2

该层的输入是"MHA的输出(output)+MHA(input_embedding)的输入，去做一个Layer normalization"得到的output(batch_size, seq_len, d_model)，经过如下代码:

self.linear1 = nn.Linear(d_model, d_ff) # 这里的d_ff在原文中设为了 2048
self.linear2 = nn.Linear(d_ff, d_model)
output = self.relu(self.linear1(inputs))
 # |output| : (batch_size, seq_len, d_ff)
output = self.linear2(output)
 # |output| : (batch_size, seq_len, d_model)

得到了相同规格的output

(4) FFN的输出+FFN的输入，去做一个Layer normalization

FFN的输出就是上面的 output (batch_size, seq_len, d_model), FFN的输入是"MHA的输出(output)+MHA(input_embedding)的输入，去做一个Layer normalization"得到的相同规格的结果。
依旧做相同的layer normalization。

self.layernorm2 = nn.LayerNorm(d_model, eps=1e-6)
在ffn_outputs = self.layernorm2(attn_outputs + ffn_outputs)

回到“TransformerEncoder” 类

已知单层的encoder layer的输出，依然是 batch_size, seq_len, d_model大小的，我们可以继续用这个输出，以及不变的attn_pad_mask，输入到下一层的encoder layer中。这个"n_layers"你可以自己定。

(5) 最后套一个线性层、激活层，用作任务的输出(不一定)

假设batch_size是句子的个数，seq_len是每个句子的长度，通过上述的一系列操作，我们最终的output已经包含了每个句子内部，单词之间的全局的关系+序列关系，可用作其它任务。(之前做句子的情感分类，会套一个输出大小为2维的线性层…)

模型回顾与总结

模型一共用到了五个类：

从"TransformerEncoder"开始，到"EncoderLayer",而它又用到了其它三个类。
写下此篇，算是跟transformer encoder部分的做了一个告别。但这个"告别"并不彻底，对于Transformer而言，我还有很多不懂的地方

• Layer Normalization的原理与作用
• Residual Dropout的含义
• Multi-Head的用处
• Decoder的原理与代码实现
• Attention的源头与变种
• Transformer的变种
• Bert的原理与使用
• …

完整代码

https://github.com/beiweixiaoxu/transformerencoder

补充: pytorch自带的transformer

才发现pytorch官网有transformers，试着对代码做一些解读。

PositionalEncoding

class PositionalEncoding(nn.Module):

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

        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).transpose(0, 1)
        self.register_buffer('pe', pe)

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

那个register_buffer是保证pe不被作为参数

attn_pad_mask

这个需要我们自己写，在padding的位置设置元素=True，然后传入就行了：

output = self.transformer_encoder(src, self.src_mask, "here for attn_pad_mask")