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. 精读

2.1 modeling.py

先看一下最主要的函数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：和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]
      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_tgt_mask：

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,
                             dtype=tf_float) # [qlen, mlen+qlen, 1, 1]
    else:
      non_tgt_mask = None复制代码