动手学Transformer

动手实现Transformer，所有代码基于tensorflow2.0，配合illustrated-transformer更香。

• 模型架构
  • Encoder+Decoder
    • Encoder
    • Decoder
    • Attention
  • Add&Norm
  • FeedForward
  • Embedding
  • Position Encoding

模型架构

transformer使用经典的编码器-解码器框架，编码器接受一个输入序列 \((x_1,…,x_n)\)，经过Embedding转化为词向量，和位置编码相加作为Encoder的输入，在一顿操作后输入被映射到\(z=(z_1,…,z_n)\)，Decoder基于\(z\)在一顿操作后生成输出序列\((y_1,…,y_m)\)。

看图说话

• 左边是Encoder，输入为词ID序列，对应形状\([batch\ size,max\ input\ sentense\ length]\)，如果embedding维度设置为512，输出形状为\([batch\ size, max\ input\ sentence\ length, 512]\),\(Nx\)表示将Encoder模块堆叠\(N\)次(论文中\(N=6\))
• 右边是Decoder，训练阶段，Decoder输入包括目标句子的词ID序列和最后一个Encoder部分的输出，测试阶段，Decoder的输入为上一次输出的词。Decoder同样被堆叠\(N\)次，最后一个Encoder的输出被接到每一个Decoder块的输入。Decoder输出下一个词的概率，输出形状为\([batch\ size, max\ output\ sentence\ length, \ vocabulary\ length])\) 。

我们先盖房子在装修

class Transformer(tf.keras.Model):
    '''
    Transformer架构，Encoder-Decoder；softmax
    params：
    num_layers：堆叠层数
    dim_model：embedding 维度
    num_heads：multihead attention
    dim_ff：FeedForWard 维度
    input_vocab_size：输入词典大小
    target_vocab_size：输出词典大小
    rate:dropout rate
    '''
    def __init__(self,num_layers, dim_model, num_heads, dim_ff, input_vocab_size,target_vocab_size, rate=0.1):
        super(Transformer, self).__init__()
        self.encoder = Encoder(num_layers, dim_model, num_heads,dim_ff, input_vocab_size, rate)#Encoder
        self.decoder = Decoder(num_layers, dim_model, num_heads,dim_ff, target_vocab_size, rate)#Decoder
        self.output_layer = tf.keras.layers.Dense(target_vocab_size)

    def call(self, inputs, targets, training, enc_padding_mask, look_ahead_mask, dec_padding_mask):
        encoder_output = self.encoder(inputs, training, enc_padding_mask)  # (batch_size, inputs_seq_len, d_model)
        decoder_output, attention_weights = self.decoder(targets, encoder_output, training, look_ahead_mask, dec_padding_mask)
        output = self.output_layer(decoder_output)  # (batch_size, tar_seq_len, target_vocab_size)
        return output, attention_weights

• Encoder

Encoder接受输入token的embedding和位置编码，经过N次Encoder layer 堆叠，代码如下所示

class Encoder(tf.keras.layers.Layer):
    '''
    Encoder 部分,input embedding ；Encoder layer stack
    '''
    def __init__(self, num_layers, dim_model, num_heads,dim_ff, input_vocab_size, rate=0.1):
        super(Encoder, self).__init__()

        self.dim_model = dim_model
        self.num_layers = num_layers

        self.embedding = tf.keras.layers.Embedding(input_vocab_size, self.dim_model)#输入Embedding
        self.pos_encoding = positional_encoding(input_vocab_size, self.dim_model)#位置编码

        self.enc_layers = [EncoderLayer(dim_model, num_heads, dim_ff, rate) for _ in range(num_layers)]#创建Encoder layer

        self.dropout = tf.keras.layers.Dropout(rate)

    def call(self, x, training, mask):

        seq_len = tf.shape(x)[1]

        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, dim_model)
    #     x *= tf.math.sqrt(tf.cast(self.dim_model, tf.float32))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for layer in self.enc_layers:
            x = layer(x, training, mask)

        return x  # (batch_size, input_seq_len, d_model)

单个Encoder layer 有两个子层： attention层和point wise feed forward network.

class EncoderLayer(tf.keras.layers.Layer):
    '''
    Encoder layer:
    multihead attention；add&layer norm；FeedForward；add&layer norm
    '''
    def __init__(self, dim_model, num_heads, dim_ff, rate=0.1):
        super(EncoderLayer, self).__init__()

        self.mha = MultiHeadAttention(dim_model, num_heads)
        self.ffn = point_wise_feed_forward_network(dim_model, dim_ff)

        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)

    def call(self, x, training, mask):

        attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)

        ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)
        ffn_output = self.dropout2(ffn_output, training=training)
        out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)

        return out2

• Decoder

Decoder和Encoder类似，由N个Decoder layer 堆叠而成， 接受输入有三部分：token embedding、位置编码、 最后一个Encoder的输出

class Decoder(tf.keras.layers.Layer):
    '''
    Decoder 部分：
    input embedding;Decoder layer stack;
    '''
    def __init__(self,num_layers,dim_model, num_heads, dim_ff, target_vocab_size, rate=0.1):
        super(Decoder, self).__init__()

        self.dim_model = dim_model
        self.num_layers = num_layers

        self.embedding = tf.keras.layers.Embedding(target_vocab_size, self.dim_model)
        self.pos_encoding = positional_encoding(target_vocab_size, self.dim_model)

        self.dec_layers = [DecoderLayer(dim_model, num_heads, dim_ff, rate) for _ in range(num_layers)]#创建Decoder layer
        self.dropout = tf.keras.layers.Dropout(rate)

    def call(self,x,enc_output,training,look_ahead_mask,padding_mask):
        seq_len = tf.shape(x)[1]
        attention_weights = {}

        x = self.embedding(x)# (batch_size, target_seq_len, d_model)
        x += self.pos_encoding[:,:seq_len,:]

        for i,layer in enumerate(self.dec_layers):
            x,block1,block2 = layer(x,enc_output,training,look_ahead_mask, padding_mask)
            attention_weights['decoder_layer{}_block1'.format(i+1)] = block1
            attention_weights['decoder_layer{}_block2'.format(i+1)] = block2

        # x.shape == (batch_size, target_seq_len, d_model)
        return x, attention_weights

单个Decoder层有三个子层：masked attention层、attention层和point wise feed forward network. masked attention层会掩盖掉序列中还没看到的位置，attention层以target token 为query，Encoder输出为 key和value

class DecoderLayer(tf.keras.layers.Layer):
    '''
    Decoder layer:
    masked multihead attention;add&norm;multihead attention;add&norm;FeedForward;add&norm
    '''
    def __init__(self, dim_model, num_heads,dim_ff,rate=0.1):
        super(DecoderLayer,self).__init__()

        self.mask_mha = MultiHeadAttention(dim_model, num_heads)
        self.mha = MultiHeadAttention(dim_model, num_heads)

        self.ffn = point_wise_feed_forward_network(dim_model, dim_ff)

        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)
        self.dropout3 = tf.keras.layers.Dropout(rate)

    def call(self,x,enc_output, training, look_ahead_mask, padding_mask):
        #masked multi-head attention
        mask_attn_output, attn_weights_block1 = self.mask_mha(x,x,x,look_ahead_mask)
        mask_attn_output = self.dropout1(mask_attn_output,training)
        out1 = self.layernorm1(x + mask_attn_output)

        #multi-head attention
        attn_output,attn_weights_block2 = self.mha(out1,enc_output,enc_output,padding_mask)
        attn_ouput = self.dropout2(attn_output, training=training)
        out2 = self.layernorm2(attn_output + out1)  # (batch_size, target_seq_len, d_model)

        # feed-forward
        ffn_output = self.ffn(out2)  # (batch_size, input_seq_len, d_model)
        ffn_output = self.dropout2(ffn_output, training=training)
        out3 = self.layernorm3(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)

        return out3, attn_weights_block1, attn_weights_block2

• Attention

点积attention

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

def Attention(query,key,value,mask=None):
    qk = tf.matmul(query,key,transpose_b=True)#对key进行转置

    dk = tf.cast(tf.shape(key)[-1],tf.float32)
    scale = qk / tf.math.sqrt(dk)

    if mask is not None:
        scale = scale + (mask*1e-9)#le-9接近于负无穷，对应softmax输出就为0

    score = tf.nn.softmax(scale,axis=-1)
    output = tf.matmul(score,value)

    return output,score

• Multihead attention

将输入划分为多个head，对应维度变化\((batch\_size,seq\_len,dim\_model) \rightarrow (batch\_size,head,seq\_len,d_k)\), 每个head操作不变，将多个head结果拼接经过权重矩阵\(W^O\)输出

\[MultiHead(Q,K,V)=Concat(head_1,...,head_h)W^O \\
where ~~~~~~ head_i=Attention(QW^Q_i,KW^K_i,VW^V_i)
\]

class MultiHeadAttention(tf.keras.layers.Layer):
    '''
    Multihead attention
    '''
    def __init__(self,dim_model,num_heads):
        super(MultiHeadAttention,self).__init__()
        self.num_heads = num_heads
        self.dim_model = dim_model

        assert dim_model%self.num_heads == 0

        self.d_k = dim_model//num_heads

        self.wq = tf.keras.layers.Dense(dim_model)
        self.wk = tf.keras.layers.Dense(dim_model)
        self.wv = tf.keras.layers.Dense(dim_model)

        self.wo = tf.keras.layers.Dense(dim_model)

    def split_heads(self,x,batch_size):
        #split 输入到多个head，(batch_size,seq_len,dim_model)->(batch_size,head,seq_len,self.d_k)
        x = tf.reshape(x,(batch_size, -1, self.num_heads, self.d_k))
        return tf.transpose(x,perm=[0,2,1,3])#重排输出维度

    def call(self,q,k,v,mask):
        batch_size = tf.shape(q)[0]

        q = self.wq(q)  # (batch_size, seq_len, dim_model)
        k = self.wk(k)
        v = self.wv(v)  

        q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, d_k)
        k = self.split_heads(k, batch_size)
        v = self.split_heads(v, batch_size)  

        attention,attention_weights = Attention(q,k,v,mask)

        attention = tf.transpose(attention,perm=[0,2,1,3])
        concat_attention = tf.reshape(attention,(batch_size,-1,self.dim_model))

        return self.wo(concat_attention),attention_weights

• Masking

在解码时，为了保证预测当前词不会看到后面的词，需要将后面的词掩码掉。如下图所示，上三角元素为1，表示该位置需要被mask

#mask掉后面的词，保证预测前面词时看不到后面的
def look_ahead_mask(size):
    "Mask out subsequent positions."
    mask = tf.linalg.band_part(tf.ones((size, size)), 0, -1)#保留上三角,1表示需要被mask
    return mask

plt.figure(figsize=(5,5))
plt.imshow(look_ahead_mask(20))
None

#序列填充mask，对padding位置输出1表示被填充
def padding_mask(seq):
  seq = tf.cast(tf.math.equal(seq, 0), tf.float32)

  # add extra dimensions so that we can add the padding
  # to the attention logits.
  return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)

• Feed-Forward network

接在attention子层后面，用于增强非线性表示能力

\[FFN(x)=max(0,xW1+b1)W2+b2
\]

def point_wise_feed_forward_network(dim_model, dim_ff):
    model = tf.keras.Sequential()
    model.add(tf.keras.layers.Dense(dim_ff, activation='relu'))
    model.add(tf.keras.layers.Dense(dim_model))
    return model

• position encoding

位置编码是影响Transformer长文本建模的主要因素，由于input embedding 中不包含词的位置信息，需要加入位置编码表示词在句子中的位置信息。

为了表示词的绝对位置信息和相对位置信息，本文中使用了正余弦函数

\[PE_{(pos,2i)}=sin(pos/10000^{2i/dmodel})\\PE_{(pos,2i+1)}=cos(pos/10000^{2i/dmodel})
\]

其中 \(pos\) 表示位置，\(i\) 表示维度，关于为什么使用正余弦作为位置编码，可以参考这里。

def positional_encoding(maxlen,dim_model):
    position = np.arange(0, maxlen)[:,np.newaxis]
    div = np.array([1/np.power(10000.0,((2*i)/dim_model)) for i in range(dim_model)])[np.newaxis,:]
#     div = np.exp(np.arange(0, dim_model, 2) * -(np.log(10000.0) / dim_model))

    sine = np.sin(position*div[:,0::2])
    cons = np.cos(position*div[:,1::2])

    pos_encoding = np.concatenate([sine,cons], axis=-1)[np.newaxis, :]

    return tf.cast(pos_encoding, dtype=tf.float32)

pos_encoding = positional_encoding(50, 512)
print (pos_encoding.shape)

plt.pcolormesh(pos_encoding[0], cmap='RdBu')
plt.xlabel('dim')
plt.xlim((0, 512))
plt.ylabel('Position')
plt.colorbar()
plt.show()

transformer的基本概念就先学到这里了

reference:

Illustrated-transformer. https://jalammar.github.io/illustrated-transformer/

Transformer_implementation_and_application. https://github.com/yuanxiaosc/Transformer_implementation_and_application

The Annotated Transformer. http://nlp.seas.harvard.edu/2018/04/03/attention.html#full-model