TF XLNet源码解读
大概读一下XLNet源码，边读边写，有问题希望和大家交流
1. 概述
1.1 文件结构
- xxx_utils.py：各种数据预处理、模型加载等辅助函数
- modeling.py：transformer-xl、tow stream attention实现
- xlnet.py：更高层的XLNetModel类，封装modeling的transformer
- function_builder.py：各种用于pretrain和finetune的loss function
- train_xxx.py：XLNet预训练
- run_xxx.py：XLNet精调和评估
总体的依赖顺序就是：
- Pretrain: train_xxx.py -> function_builder.py -> modeling.py -> xxx_utils.py
- Finetune: run_xxx.py -> function_builder.py -> modeling.py -> xxx_utils.py
最精华且难啃的部分就是modeling.py，其他的看一下就差不错了，主要是一起读一下这个文件，之后其他的再慢慢加
2. 精读
这里主要读一下modeling.py脚本，其他的有基础的同学看一看应该就能明白～
2.1 先看一下最主要的函数transformer_xl，代码太多就不全贴了，挑一些重点的
- 输入参数
- mems：这个存了前mem_len个batch的信息，estimator每计算一个batch会更新一次，都存在TrainSpec里
- perm_mask：[i, j, k]表示在第k个batch，i和j计算attention(0)、不计算(1)，因为要加上之前的mems计算，所以会多出k维度和各个batch对齐
- target_mapping：因为理论上把token都permute了，所以可能先预测4再预测2，所以在预测i=0（第一个4）时要把实际的位置4给mask掉。这里作者说“in batch k”感觉有些不对，这个应该只针对当前的batch，k应该表示的是batch里的第k个
- inp_q：没理解错的话，1的token相当于BERT的[MASK]，如果是None的话就不进行PLM任务
- untier：是否统一attention计算中的bias。之前BERT对于multi-head的投影都是直接用dense，这里projection矩阵和bias矩阵是分开的，而且untie_r=False时所有layer的bias都一样
- clamp_len：限制relative的长度
- bias：这里有三种，论文中称为head specific bias vector，我觉得应该是为了增强拟合能力。有content attention的r_w_bias，position attention的r_r_bias，segment attention的r_s_bias，在rel_attn_core函数中看的比较明白：
def rel_attn_core(q_head, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat,
r_w_bias, r_r_bias, r_s_bias, attn_mask, dropatt, is_training,
scale):
"""Core relative positional attention operations."""
# content based attention score
ac = tf.einsum('ibnd,jbnd->ijbn', q_head + r_w_bias, k_head_h)
# position based attention score
bd = tf.einsum('ibnd,jbnd->ijbn', q_head + r_r_bias, k_head_r)
bd = rel_shift(bd, klen=tf.shape(ac)[1])
# segment based attention score
if seg_mat is None:
ef = 0
else:
ef = tf.einsum('ibnd,snd->ibns', q_head + r_s_bias, seg_embed)
ef = tf.einsum('ijbs,ibns->ijbn', seg_mat, ef)
# merge attention scores and perform masking
attn_score = (ac + bd + ef) * scale
# more ...- attn_mask：计算query stream的mask，需要和attention_score保持一致，转换为4维
if data_mask is not None: # [1, len, bsz] + [len, len, bsz] = [qlen, qlen, bsz]
# all mems can be attended to
mems_mask = tf.zeros([tf.shape(data_mask)[0], mlen, bsz],
dtype=tf_float) # [qlen, mlen, bsz] 全0矩阵
data_mask = tf.concat([mems_mask, data_mask], 1) # [qlen, mlen+qlen, bsz]
if attn_mask is None:
attn_mask = data_mask[:, :, :, None] # [qlen, mlen+qlen, bsz, 1]
else:
attn_mask += data_mask[:, :, :, None]- non_tgtmask：计算content stream的mask，也是finetune阶段的mask，由于content stream是可以看到自己的，所以初始化时是 -tf.eye，这样和attnmask相加时就把对自己的mask减去了（attn_mask里是1，加上non_tgt_mask相当于减去1）
if attn_mask is not None:
non_tgt_mask = -tf.eye(qlen, dtype=tf_float) # [qlen, qlen]单位矩阵
non_tgt_mask = tf.concat([tf.zeros([qlen, mlen], dtype=tf_float), # [qlen, mlen+qlen]
non_tgt_mask], axis=-1)
non_tgt_mask = tf.cast((attn_mask + non_tgt_mask[:, :, None, None]) > 0, #减去对自己的mask
dtype=tf_float) # [qlen, mlen+qlen, 1, 1]
else:
non_tgt_mask = None- word_emb：这里有两个，word_emb_k就是根据词表id的随机初始化词向量，也作为content stream attention的初始值（能看到每个token的emb）。word_emb_q主要是为了生成query stream的初始值，因为要去掉被mask的word_emb，借助了一下inp_q，另外给所有mask的token放了一个mask_emb，就像BERT一样，[MASK]本身也有自己的emb。
if inp_q is not None:
with tf.variable_scope('mask_emb'):
mask_emb = tf.get_variable('mask_emb', [1, 1, d_model], dtype=tf_float) # [1, 1, d_model]
if target_mapping is not None:
word_emb_q = tf.tile(mask_emb, [tf.shape(target_mapping)[0], bsz, 1]) # [num_predict, bsz, d_model]
else:
inp_q_ext = inp_q[:, :, None] # [len, bsz， 1]
# word_emb_q = 用mask_emb代替被mask掉的word_emb_k
word_emb_q = inp_q_ext * mask_emb + (1 - inp_q_ext) * word_emb_k # [len, bsz, d_model]
output_h = tf.layers.dropout(word_emb_k, dropout, training=is_training)
if inp_q is not None:
output_g = tf.layers.dropout(word_emb_q, dropout, training=is_training)- seg_embed：相对seg_embed，不管多少个segment，只需要判断0（是同一个segment），1（不是），维度的设置方便之后直接计算attention，至于为什么有n_layer我有点不懂？？？
- seg_mat：一个one hot矩阵，0（是同一个segment），1（不是），这里运算的操作有些nb，反正大家知道最后结果是什么样子就可以
# `1` indicates not in the same segment [qlen x klen x bsz]
seg_mat = tf.cast(
tf.logical_not(tf.equal(seg_id[:, None], cat_ids[None, :])),
tf.int32) # [qlen, 1, bsz] [1, qlen+mlen, bsz] => [qlen x klen x bsz]
seg_mat = tf.one_hot(seg_mat, 2, dtype=tf_float)最先处理的几个矩阵都初始化好了，就开始进入n_layer的transformer-xl阶段，其中pretrain的时候需要用two_stream_rel_attn，预训练和评估时使用rel_multihead_attn。还有一个点，就是reuse的设置，在pretrain的时候是True，finetune时候是False，这样预训练时两个stream共享content stream attention之后的positionwise_ffn的权重。（感谢 @谢暄 同学的评论）
2.2 接下来先啃一下rel_multihead_attn
这里与传统attention的区别主要是要分开计算position和segment的attention，最终加起来（实现在rel_attn_core里）。加了一些维度的注释：
# h - [len, bsz, d_model]
# proj_w - [d_model, n_head, d_head]
# q_head_h - [len, bsz, n_head, d_head]
# content heads
q_head_h = head_projection(
h, d_model, n_head, d_head, kernel_initializer, 'q')
k_head_h = head_projection(
cat, d_model, n_head, d_head, kernel_initializer, 'k')
v_head_h = head_projection(
cat, d_model, n_head, d_head, kernel_initializer, 'v')
# r - [len, bsz, d_model]
# positional heads
k_head_r = head_projection(
r, d_model, n_head, d_head, kernel_initializer, 'r')
# core attention ops
# attn_vec - [len, len, bsz, n_head]
# 分别计算content、position、segment的attention score
attn_vec = rel_attn_core(
q_head_h, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
r_r_bias, r_s_bias, attn_mask, dropatt, is_training, scale)
# post processing: residual+layernorm
# output - [len, bsz, d_model]
output = post_attention(h, attn_vec, d_model, n_head, d_head, dropout,
is_training, kernel_initializer)- two_stream_rel_attn
这里其实就是比rel_multihead_attn多了一个query stream，由传进来的g单独计算出output_g，原理都是一样的。如此聪明的大家肯定都能举一反三～
2.3 关于relative position的一系列操作
首先是生成embedding，主要逻辑在relative_positional_encoding，先是生成了前后向的相对位置，然后再clip，之后通过positional_embedding生成最终embedding，这里大致操作我注释了下维度，但是为什么这么操作不太懂？？？
def positional_embedding(pos_seq, inv_freq, bsz=None):
sinusoid_inp = tf.einsum('i,d->id', pos_seq, inv_freq) # [len, d_model//2]
pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], -1)
pos_emb = pos_emb[:, None, :] # [len, 1, d_model]
if bsz is not None: # bsz = orig_bsz//2
pos_emb = tf.tile(pos_emb, [1, bsz, 1]) # [len, bsz, d_model]
return pos_emb之后就是在计算rel_attn_core的时候对计算好的attention进行了rel_shift，这里我真的懵。。先贴上一些注释，休息过来再读一下论文？？？
def rel_shift(x, klen=-1):
"""perform relative shift to form the relative attention score."""
x_size = tf.shape(x) # [qlen, klen, bsz, n_head]
x = tf.reshape(x, [x_size[1], x_size[0], x_size[2], x_size[3]]) # [klen, qlen, bsz, n_head]
x = tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1]) # [klen-1, qlen, bsz, n_head]
x = tf.reshape(x, [x_size[0], x_size[1] - 1, x_size[2], x_size[3]]) # [qlen, klen-1, bsz, n_head]
x = tf.slice(x, [0, 0, 0, 0], [-1, klen, -1, -1]) # ???目前先到这里，希望也看过源码的同学一起讨论一下
编辑于 2020-09-09 13:14
