Transformer最详细的原理加代码解读_transformer原理与代码精讲

Transformer原理

1. motivation

​ 为了解决seq2seq的问题，之前一般都是使用RNN模型进行求解。RNN的一大劣势就是无法进行并行化计算，比如要想输出$b^4$就必须要先获得$a^1$到$a^4$才行。而接下来就有学者想采把CNN用来取代RNN，每个小三角形都是一个filter，但是问题是如下图所示每个小三角仅能考虑到很少的一部分输入，但是我们可以通过叠多层的CNN，则上层的filter就可以考虑到比较多的语句，如下所示蓝色的filter可以看到$b^1$到$b^3$，而$b^1$到$b^3$是由$a^1$到$a^4$决定的，相当于蓝色的filter已经看到了整个句子中所有的内容。但是这里有个问题是你需要叠很多层CNN才能够看到整个句子，那如果你想第二层就看到整个句子就会比较困难，所以有一个新的想法就是self-attention。

​ self-attention就是做了这样一件事情，输入和输出和RNN一样都是seq，然后$b^1$到$b^4$可以看到整个输入seq，但是$b^1$到$b^4$却又可以被并行计算。

2. self-attention原理

2.1 self-attention计算过程

​ 如下所示，首先input是$x^1$到$x^4$（一个sequence），每个input先通过一个embedding（乘以一个W矩阵）变成$a^1$到$a^4$，然后把他们丢进self-attention layer。self-attention里面做的是，先把每个input都乘上三个不同的matrix产生三个不同的vector，这三个不同的vector我们分别命名为q、k、v。q代表query（to match others），把每一个input a都乘上某一个matrix $W^q$，然后就得到了$q^1$到$q^4$，这些东西叫做query。用同样的方法把input乘上$W^k$就得到$k^1$到$k^4$，这个k叫做key，用来被querymatch的。v是指被抽取出来的信息，获取方式也是同理。

​ 接下来要做的是拿每一个q对每一个k做attention，我们先把$q^1$拿出来对$k^1$做attention得到$a_{1,1}$，那接下来拿$q^1$对$k^2$做attention得到$a_{1,2}$，拿$q^1$对$k^3$做attention得到$a_{1,3}$，拿$q^1$对$k^4$做attention得到$a_{1,4}$。那这个a是如何计算的呢？其实就是如下所示的Scaled Dot-Product Attention，为什么要除以根号d呢？因为q和k的dot product数值会随着d（q和k的维度）的增大而增大，很直观就是q和k的维度越大，相乘之后再相加的数就越多。为什么是除以根号d而不是别的呢？原文中有个注脚有写，但是具体的目前还是不清楚，有人说根号d是q和k的标准差，但是无从考证。

​ 接下来就是做一个softmax，得到$a_{1,1}$到$a_{1,4}$ head.

​ 有了$\widehat{a_{1,1}}$ 到$\widehat{a_{1,4}}$ 之后就那每一个a去和v相乘再做weight sum得到$b^1$，你会发现产生$b^1$的时候用到了整个sequence，因为$b^1$用到了$v^1$到$v^4$，而$v^1$到$v^4$使用了全部的$a^1$到$a^4$。

​ $b^2$的计算过程也和前面的类似，这里就不展开讲了。

​ 可以看到，$b^1$到$b^4$是可以被并行计算的。

​

2.2 self-attention矩阵计算过程

​ 下面可以把前面介绍的那些过程用矩阵运算来进行计算。把前面的经过embedding转化的$a^1$到$a^4$乘以可学习矩阵$W^q$就得到了$q^1$到$q^4$，同理$k^1$到$k^4$以及$v^1$到$v^4$也是如此计算。

​ 这里为了方便计算忽略分母的根号d，$k^1$做一个转置然后乘以$q^1$就得到了$a_{1,1}$，那我们把$k^1$到$k^4$叠在一起变成一个matrix，直接把这个matrix乘以$q^1$得到的结果是一个向量，这个向量就是$a_{1,1}$到$a_{1,4}$，$a_{1,1}$到$a_{1,4}$的计算也是可以并行的。

​ 同理$q^2$拿出来和$k^1$到$k^4$相乘就得到了$a_{2,1}$到$a_{2,4}$，剩下的$q^3$和$q^4$也是同理。这个整个计算attention的过程其实就相当于把K矩阵做一个转置，直接乘上Q，即得到了attention A，每一个input两两之间都有attention，如果这边有四个input那得到的attention就是4*4，如果输入sequence长度是n则得到的attention矩阵就是$n^2$。接下来把A做一个softmax就可以。

​ 接下来就是把$v_1$到$v_4$根据$\hat{a}$做weight sum。我们就是把$v_1$和$\widehat{a_{1,1}}$到$\widehat{a_{1,4}}$对应相乘再相加得到$b_1$，同理得到$b_2$到$b_4$，最后把$b_1$到$b_4$串起来就得到了最终的结果O

​ 最后让我们来看一下整体的矩阵运算，从输入到输出就是一堆矩阵乘法，而矩阵乘法是很轻易地可以用GPU加速的。

2.3 multi-head机制

​ 这里就用2head的情况进行举例，每一个$a^i$都会得到$q^i$、$k^i$和$v^i$，在multi-head机制下，你会继续把$q^i$乘上两个不同的matrix（$W^{q,1}$和$W^{q,2}$）得到$q^{i,1}$和$q^{i,2}$，同理可以产生$k^{i,1}$和$k^{i,2}$，以及$v^{i,1}$和$v^{i,2}$。但是这里$q^{i,1}$只会和$k^{i,1}$和$k^{j,1}$做attention最后计算出$b^{i,1}$。

​ 和上面原理类似，你也可以计算出$b^{i,2}$，然后我们会把他们concat起来。

​ 那如果concat之后的维度过大，你可以乘以一个矩阵$W^o$使其降维，最终就得到了$b^i$。那采用multi-head有什么好处吗？其实不同的head他们关注的点不一样，举例来说有的head可能想要看的就是local的信息，有的head想要看的是比较长时间的信息，就是从不同的角度看问题。

2.4 Positional Encoding

​ 前面self-attention过程中你会发现，input的顺序是不重要的。因为他做的事情就是对每一个input的内容都做self-attention，每一个时间点来说跟它是邻居还是在远处对它来说都是一样的，在self-attention心里并没有所谓位置的概念。a点乘b还是b点乘a对它来说好像是一样的，显然我们不希望这样，我们希望能够把input sequence的顺序考虑进来到self-attention里面去。

​ 最初的transformer论文中的做法是人为设计的一个表示位置的矩阵，每一个位置i都有一个向量$e^i$，这个向量就可以表示该词的位置信息。然后这个位置向量$e^i$(维度和$a^i$相同)直接加到原始的$a^i$上就融合了位置信息了，那为什么是相加而不是concat呢？这里李宏毅老师给出了一种自己的解释，我们把$x^i$再填上一个维度的one-hot向量叫$p^i$，就是说第i维是1，其他都是0。看到这个one-hot向量你就知道这个词落在哪个位置，然后把$x^i$和$p^i$做concat，然后把它乘上一个W生成embedding然后做transformer，那我们可以把这个W拆成两个矩阵$W^I$和$W^P$，这样相当于于把$x^i$和$p^i$串起来乘以W转换为把$W^I$乘以$x^i$再加上$W^P$乘以$p^i$。而这两个部分恰好分别是$a^i$和$e^i$，因此直接把$a^i$和$e^i$相加并没有什么奇怪。

​ 这里又有一个新的疑问就是$W^P$如何确定，论文中给出了就是长成如下的样子（可用正弦和余弦表示），后面的bert直接就设为一个科学习的参数矩阵了。

3. Seq2seq with Attention

​ transformer整体结构如下：

​ 具体的细节如下图所示，这里讲一下batch norm和layer norm的区别，假设我们有一个batch的数据，batchsize=4，BN是对同一个batch里面不同的data的同样的dimension做norm，我们希望同一个batch里面同一个dimension的均值等于0，方差等于1。LN不需要考虑batch的，LN就是给你data我们希望各个不同dimension的均值是0，方差是1。

4. 源码解读

​ github：https://github.com/Kyubyong/transformer

4.1 数据生成

​ 数据生成的代码在根目录下的train.py中，使用的是get_batch这个函数，分别生成train_batches和eval_batches，后面生成一个迭代器iter，迭代的获取训练数据和标签。

train_batches, num_train_batches, num_train_samples = get_batch(hp.train1, hp.train2,
                                             hp.maxlen1, hp.maxlen2,
                                             hp.vocab, hp.batch_size,
                                             shuffle=True)
eval_batches, num_eval_batches, num_eval_samples = get_batch(hp.eval1, hp.eval2,
                                             100000, 100000,
                                             hp.vocab, hp.batch_size,
                                             shuffle=False)

# create a iterator of the correct shape and type
iter = tf.data.Iterator.from_structure(train_batches.output_types, train_batches.output_shapes)
xs, ys = iter.get_next()

train_init_op = iter.make_initializer(train_batches)
eval_init_op = iter.make_initializer(eval_batches)

logging.info("# Load model")
m = Transformer(hp)    # 主模型
loss, train_op, global_step, train_summaries = m.train(xs, ys) # train这个模型
y_hat, eval_summaries = m.eval(xs, ys)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

4.2 transformer类

​ transformer类是模型中最重要的类，下面将介绍encoder、decoder、train、eval部分的代码。

4.2.1 encoder

​ encoder的输入是训练数据xs，先经过查找表生成embedding，然后加上位置编码信息后经过dropout层。中间就是多个multi-head attention模块，最后再经过一个feed forward。

def encode(self, xs, training=True):
        '''
        Returns
        memory: encoder outputs. (N, T1, d_model)
        '''
        with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
            x, seqlens, sents1 = xs

            # src_masks
            src_masks = tf.math.equal(x, 0) # (N, T1)

            # embedding
            enc = tf.nn.embedding_lookup(self.embeddings, x) # (N, T1, d_model)
            enc *= self.hp.d_model**0.5 # scale

            enc += positional_encoding(enc, self.hp.maxlen1)
            enc = tf.layers.dropout(enc, self.hp.dropout_rate, training=training)

            ## Blocks
            for i in range(self.hp.num_blocks):
                with tf.variable_scope("num_blocks_{}".format(i), reuse=tf.AUTO_REUSE):
                    # self-attention
                    enc = multihead_attention(queries=enc,
                                              keys=enc,
                                              values=enc,
                                              key_masks=src_masks,
                                              num_heads=self.hp.num_heads,
                                              dropout_rate=self.hp.dropout_rate,
                                              training=training,
                                              causality=False)
                    # feed forward
                    enc = ff(enc, num_units=[self.hp.d_ff, self.hp.d_model])
        memory = enc
        return memory, sents1, src_masks
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

​ 具体的positional_encoding代码如下，分奇偶使用sin和cos进行位置编码：

def positional_encoding(inputs,
                        maxlen,
                        masking=True,
                        scope="positional_encoding"):
    '''Sinusoidal Positional_Encoding. See 3.5
    inputs: 3d tensor. (N, T, E)
    maxlen: scalar. Must be >= T
    masking: Boolean. If True, padding positions are set to zeros.
    scope: Optional scope for `variable_scope`.

    returns
    3d tensor that has the same shape as inputs.
    '''

    E = inputs.get_shape().as_list()[-1] # static
    N, T = tf.shape(inputs)[0], tf.shape(inputs)[1] # dynamic
    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
        # position indices
        position_ind = tf.tile(tf.expand_dims(tf.range(T), 0), [N, 1]) # (N, T)

        # First part of the PE function: sin and cos argument
        position_enc = np.array([
            [pos / np.power(10000, (i-i%2)/E) for i in range(E)]
            for pos in range(maxlen)])

        # Second part, apply the cosine to even columns and sin to odds.
        position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])  # dim 2i
        position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])  # dim 2i+1
        position_enc = tf.convert_to_tensor(position_enc, tf.float32) # (maxlen, E)

        # lookup
        outputs = tf.nn.embedding_lookup(position_enc, position_ind)

        # masks
        if masking:
            outputs = tf.where(tf.equal(inputs, 0), inputs, outputs)

        return tf.to_float(outputs)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

​ multihead_attention的具体代码如下，先当于把Q、K、V先做线性投影，然后做split再concat，最后Q_ 、K_ 、V_的第一个维度（batchsize）增加了N倍，第三个维度（隐藏层维度）减小N倍。

def multihead_attention(queries, keys, values, key_masks,
                        num_heads=8, 
                        dropout_rate=0,
                        training=True,
                        causality=False,
                        scope="multihead_attention"):
    '''Applies multihead attention. See 3.2.2
    queries: A 3d tensor with shape of [N, T_q, d_model].
    keys: A 3d tensor with shape of [N, T_k, d_model].
    values: A 3d tensor with shape of [N, T_k, d_model].
    key_masks: A 2d tensor with shape of [N, key_seqlen]
    num_heads: An int. Number of heads.
    dropout_rate: A floating point number.
    training: Boolean. Controller of mechanism for dropout.
    causality: Boolean. If true, units that reference the future are masked.
    scope: Optional scope for `variable_scope`.
        
    Returns
      A 3d tensor with shape of (N, T_q, C)  
    '''
    d_model = queries.get_shape().as_list()[-1]
    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
        # Linear projections
        Q = tf.layers.dense(queries, d_model, use_bias=True) # (N, T_q, d_model)
        K = tf.layers.dense(keys, d_model, use_bias=True) # (N, T_k, d_model)
        V = tf.layers.dense(values, d_model, use_bias=True) # (N, T_k, d_model)
        
        # Split and concat
        Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0) # (h*N, T_q, d_model/h)
        K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0) # (h*N, T_k, d_model/h)
        V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0) # (h*N, T_k, d_model/h)

        # Attention
        outputs = scaled_dot_product_attention(Q_, K_, V_, key_masks, causality, dropout_rate, training)

        # Restore shape
        outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2 ) # (N, T_q, d_model)
              
        # Residual connection
        outputs += queries
              
        # Normalize
        outputs = ln(outputs)
 
    return outputs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

4.2.2 decoder

​ 下面是decoder的代码。

def decode(self, ys, memory, src_masks, training=True):
        '''
        memory: encoder outputs. (N, T1, d_model)
        src_masks: (N, T1)

        Returns
        logits: (N, T2, V). float32.
        y_hat: (N, T2). int32
        y: (N, T2). int32
        sents2: (N,). string.
        '''
        with tf.variable_scope("decoder", reuse=tf.AUTO_REUSE):
            decoder_inputs, y, seqlens, sents2 = ys

            # tgt_masks
            tgt_masks = tf.math.equal(decoder_inputs, 0)  # (N, T2)

            # embedding
            dec = tf.nn.embedding_lookup(self.embeddings, decoder_inputs)  # (N, T2, d_model)
            dec *= self.hp.d_model ** 0.5  # scale

            dec += positional_encoding(dec, self.hp.maxlen2)
            dec = tf.layers.dropout(dec, self.hp.dropout_rate, training=training)

            # Blocks
            for i in range(self.hp.num_blocks):
                with tf.variable_scope("num_blocks_{}".format(i), reuse=tf.AUTO_REUSE):
                    # Masked self-attention (Note that causality is True at this time)
                    dec = multihead_attention(queries=dec,
                                              keys=dec,
                                              values=dec,
                                              key_masks=tgt_masks,
                                              num_heads=self.hp.num_heads,
                                              dropout_rate=self.hp.dropout_rate,
                                              training=training,
                                              causality=True,
                                              scope="self_attention")

                    # Vanilla attention
                    dec = multihead_attention(queries=dec,
                                              keys=memory,
                                              values=memory,
                                              key_masks=src_masks,
                                              num_heads=self.hp.num_heads,
                                              dropout_rate=self.hp.dropout_rate,
                                              training=training,
                                              causality=False,
                                              scope="vanilla_attention")
                    ### Feed Forward
                    dec = ff(dec, num_units=[self.hp.d_ff, self.hp.d_model])

        # Final linear projection (embedding weights are shared)
        weights = tf.transpose(self.embeddings) # (d_model, vocab_size)
        logits = tf.einsum('ntd,dk->ntk', dec, weights) # (N, T2, vocab_size)
        y_hat = tf.to_int32(tf.argmax(logits, axis=-1))

        return logits, y_hat, y, sents2
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

4.2.3 train

​ 训练代码，包含encoder、decoder和参数优化更新的代码。

def train(self, xs, ys):
        '''
        Returns
        loss: scalar.
        train_op: training operation
        global_step: scalar.
        summaries: training summary node
        '''
        # forward
        memory, sents1, src_masks = self.encode(xs)
        logits, preds, y, sents2 = self.decode(ys, memory, src_masks)

        # train scheme
        y_ = label_smoothing(tf.one_hot(y, depth=self.hp.vocab_size))
        ce = tf.nn.softmax_cross_entropy_with_logits_v2(logits=logits, labels=y_)
        nonpadding = tf.to_float(tf.not_equal(y, self.token2idx["<pad>"]))  # 0: <pad>
        loss = tf.reduce_sum(ce * nonpadding) / (tf.reduce_sum(nonpadding) + 1e-7)

        global_step = tf.train.get_or_create_global_step()
        lr = noam_scheme(self.hp.lr, global_step, self.hp.warmup_steps)
        optimizer = tf.train.AdamOptimizer(lr)
        train_op = optimizer.minimize(loss, global_step=global_step)

        tf.summary.scalar('lr', lr)
        tf.summary.scalar("loss", loss)
        tf.summary.scalar("global_step", global_step)

        summaries = tf.summary.merge_all()

        return loss, train_op, global_step, summaries
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

4.2.4 eval

​ 评估代码。

def eval(self, xs, ys):
        '''Predicts autoregressively
        At inference, input ys is ignored.
        Returns
        y_hat: (N, T2)
        '''
        decoder_inputs, y, y_seqlen, sents2 = ys

        decoder_inputs = tf.ones((tf.shape(xs[0])[0], 1), tf.int32) * self.token2idx["<s>"]
        ys = (decoder_inputs, y, y_seqlen, sents2)

        memory, sents1, src_masks = self.encode(xs, False)

        logging.info("Inference graph is being built. Please be patient.")
        for _ in tqdm(range(self.hp.maxlen2)):
            logits, y_hat, y, sents2 = self.decode(ys, memory, src_masks, False)
            if tf.reduce_sum(y_hat, 1) == self.token2idx["<pad>"]: break

            _decoder_inputs = tf.concat((decoder_inputs, y_hat), 1)
            ys = (_decoder_inputs, y, y_seqlen, sents2)

        # monitor a random sample
        n = tf.random_uniform((), 0, tf.shape(y_hat)[0]-1, tf.int32)
        sent1 = sents1[n]
        pred = convert_idx_to_token_tensor(y_hat[n], self.idx2token)
        sent2 = sents2[n]

        tf.summary.text("sent1", sent1)
        tf.summary.text("pred", pred)
        tf.summary.text("sent2", sent2)
        summaries = tf.summary.merge_all()

        return y_hat, summaries
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33