[Transformer-XL]论文实现：Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context

论文：Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context
作者：Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, Ruslan Salakhutdinov
时间：2019
地址：kimiyoung/transformer-xl (github.com)

一、完整代码

有时间用tensorflow2.0实现一般，原文tensorflow1.0版本模型的代码如下：

# 完整代码在这里
import tensorflow as tf


def positional_embedding(pos_seq, inv_freq, bsz=None):
  sinusoid_inp = tf.einsum('i,j->ij', pos_seq, inv_freq)
  pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], -1)
  if bsz is not None:
    return tf.tile(pos_emb[:, None, :], [1, bsz, 1])
  else:
    return pos_emb[:, None, :]


def positionwise_FF(inp, d_model, d_inner, dropout, kernel_initializer,
                    scope='ff', is_training=True):
  output = inp
  with tf.variable_scope(scope):
    output = tf.layers.dense(inp, d_inner, activation=tf.nn.relu,
                             kernel_initializer=kernel_initializer,
                             name='layer_1')
    output = tf.layers.dropout(output, dropout, training=is_training,
                               name='drop_1')
    output = tf.layers.dense(output, d_model,
                             kernel_initializer=kernel_initializer,
                             name='layer_2')
    output = tf.layers.dropout(output, dropout, training=is_training,
                               name='drop_2')
    output = tf.contrib.layers.layer_norm(output + inp, begin_norm_axis=-1)
  return output


def rel_shift(x):
  x_size = tf.shape(x)

  x = tf.pad(x, [[0, 0], [1, 0], [0, 0], [0, 0]])
  x = tf.reshape(x, [x_size[1] + 1, x_size[0], x_size[2], x_size[3]])
  x = tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1])
  x = tf.reshape(x, x_size)

  return x


def rel_multihead_attn(w, r, r_w_bias, r_r_bias, attn_mask, mems, d_model,
                       n_head, d_head, dropout, dropatt, is_training,
                       kernel_initializer, scope='rel_attn'):
  scale = 1 / (d_head ** 0.5)
  with tf.variable_scope(scope):
    qlen = tf.shape(w)[0]
    rlen = tf.shape(r)[0]
    bsz = tf.shape(w)[1]

    cat = tf.concat([mems, w],
                    0) if mems is not None and mems.shape.ndims > 1 else w
    w_heads = tf.layers.dense(cat, 3 * n_head * d_head, use_bias=False,
                              kernel_initializer=kernel_initializer, name='qkv')
    r_head_k = tf.layers.dense(r, n_head * d_head, use_bias=False,
                               kernel_initializer=kernel_initializer, name='r')

    w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, -1)
    w_head_q = w_head_q[-qlen:]

    klen = tf.shape(w_head_k)[0]

    w_head_q = tf.reshape(w_head_q, [qlen, bsz, n_head, d_head])
    w_head_k = tf.reshape(w_head_k, [klen, bsz, n_head, d_head])
    w_head_v = tf.reshape(w_head_v, [klen, bsz, n_head, d_head])

    r_head_k = tf.reshape(r_head_k, [rlen, n_head, d_head])

    rw_head_q = w_head_q + r_w_bias
    rr_head_q = w_head_q + r_r_bias

    AC = tf.einsum('ibnd,jbnd->ijbn', rw_head_q, w_head_k)
    BD = tf.einsum('ibnd,jnd->ijbn', rr_head_q, r_head_k)
    BD = rel_shift(BD)

    attn_score = (AC + BD) * scale
    attn_mask_t = attn_mask[:, :, None, None]
    attn_score = attn_score * (1 - attn_mask_t) - 1e30 * attn_mask_t

    attn_prob = tf.nn.softmax(attn_score, 1)
    attn_prob = tf.layers.dropout(attn_prob, dropatt, training=is_training)

    attn_vec = tf.einsum('ijbn,jbnd->ibnd', attn_prob, w_head_v)
    size_t = tf.shape(attn_vec)
    attn_vec = tf.reshape(attn_vec, [size_t[0], size_t[1], n_head * d_head])

    attn_out = tf.layers.dense(attn_vec, d_model, use_bias=False,
                               kernel_initializer=kernel_initializer, name='o')
    attn_out = tf.layers.dropout(attn_out, dropout, training=is_training)

    output = tf.contrib.layers.layer_norm(attn_out + w, begin_norm_axis=-1)
  return output


def embedding_lookup(lookup_table, x, use_tpu=True):
    return tf.nn.embedding_lookup(lookup_table, x)


def mask_adaptive_embedding_lookup(x, n_token, d_embed, d_proj, cutoffs, initializer,
  return y, ret_params


def mul_adaptive_embedding_lookup(x, n_token, d_embed, d_proj, cutoffs, initializer,
  return y, ret_params


def mask_adaptive_logsoftmax(hidden, target, n_token, d_embed, d_proj, cutoffs,
  return nll


def mul_adaptive_logsoftmax(hidden, target, n_token, d_embed, d_proj, cutoffs,
  return nll


def _create_mask(qlen, mlen, same_length=False):
  return ret

def _cache_mem(curr_out, prev_mem, mem_len=None):
  return tf.stop_gradient(new_mem)


def transformer(dec_inp, target, mems, n_token, n_layer, d_model, d_embed,
                n_head, d_head, d_inner, dropout, dropatt,
                initializer, is_training, proj_initializer=None,
                mem_len=None, cutoffs=[], div_val=1, tie_projs=[],
                same_length=False, clamp_len=-1, use_tpu=True,
                input_perms=None, target_perms=None, head_target=None,
                untie_r=False, proj_same_dim=True,
                scope='transformer'):
  """
  cutoffs: a list of python int. Cutoffs for adaptive softmax.
  tie_projs: a list of python bools. Whether to tie the projections.
  use_tpu: if True, use one_hot in embedding lookup and bin-based implementation
        of adaptive softmax.
  perms: a list of tensors. Each tensor should of size [len, bsz, bin_size].
        Only used in the adaptive setting.
  """
  new_mems = []
  with tf.variable_scope(scope):
    if untie_r:
      r_w_bias = tf.get_variable('r_w_bias', [n_layer, n_head, d_head],
                               initializer=initializer)
      r_r_bias = tf.get_variable('r_r_bias', [n_layer, n_head, d_head],
                                 initializer=initializer)
    else:
      r_w_bias = tf.get_variable('r_w_bias', [n_head, d_head],
                                 initializer=initializer)
      r_r_bias = tf.get_variable('r_r_bias', [n_head, d_head],
                                 initializer=initializer)

    qlen = tf.shape(dec_inp)[0]
    mlen = tf.shape(mems[0])[0] if mems is not None else 0
    klen = mlen + qlen

    if proj_initializer is None:
      proj_initializer = initializer
    lookup_fn = (mul_adaptive_embedding_lookup if use_tpu else
                 mask_adaptive_embedding_lookup)
    embeddings, shared_params = lookup_fn(
        x=dec_inp,
        n_token=n_token,
        d_embed=d_embed,
        d_proj=d_model,
        cutoffs=cutoffs,
        initializer=initializer,
        proj_initializer=proj_initializer,
        div_val= div_val,
        perms=input_perms,
        proj_same_dim=proj_same_dim)

    attn_mask = _create_mask(qlen, mlen, same_length)

    pos_seq = tf.range(klen - 1, -1, -1.0)
    if clamp_len > 0:
      pos_seq = tf.minimum(pos_seq, clamp_len)
    inv_freq = 1 / (10000 ** (tf.range(0, d_model, 2.0) / d_model))
    pos_emb = positional_embedding(pos_seq, inv_freq)

    output = tf.layers.dropout(embeddings, dropout, training=is_training)
    pos_emb = tf.layers.dropout(pos_emb, dropout, training=is_training)

    if mems is None:
      mems = [None] * n_layer

    for i in range(n_layer):
      # cache new mems
      new_mems.append(_cache_mem(output, mems[i], mem_len))

      with tf.variable_scope('layer_{}'.format(i)):
        output = rel_multihead_attn(
            w=output,
            r=pos_emb,
            r_w_bias=r_w_bias if not untie_r else r_w_bias[i],
            r_r_bias=r_r_bias if not untie_r else r_r_bias[i],
            attn_mask=attn_mask,
            mems=mems[i],
            d_model=d_model,
            n_head=n_head,
            d_head=d_head,
            dropout=dropout,
            dropatt=dropatt,
            is_training=is_training,
            kernel_initializer=initializer)
        output = positionwise_FF(
            inp=output,
            d_model=d_model,
            d_inner=d_inner,
            dropout=dropout,
            kernel_initializer=initializer,
            is_training=is_training)

    output = tf.layers.dropout(output, dropout, training=is_training)

    logsoftmax_fn = (mul_adaptive_logsoftmax if use_tpu else
                     mask_adaptive_logsoftmax)
    loss = logsoftmax_fn(
        hidden=output,
        target=target,
        n_token=n_token,
        d_embed=d_embed,
        d_proj=d_model,
        cutoffs=cutoffs,
        params=shared_params,
        tie_projs=tie_projs,
        initializer=initializer,
        proj_initializer=proj_initializer,
        div_val=div_val,
        perms=target_perms,
        head_target=head_target,
        proj_same_dim=proj_same_dim)
    return loss, new_mems

二、论文解读

Transformer有一个学习长期依赖的潜力，但是其收到了上下文长度的限制，这里提出了一个神经网络架构Transformer-XL能够突破这一限制，而不破坏时间一致性；该模型有两个重要的组成成分：一个是分段的递归机制；一个是相对位置编码；该架构学习长期依赖性要比RNNs长80%，比普通的Transformer要长450%；其在短序列和长序列上都有更好的性能，同时在evaluation过程中比普通的Transformer快1800+倍；

2.1 Transformer-XL的介绍

首先要介绍一下什么是长期依赖性；

长期依赖性：

在序列长度无限长的情况下，序列中元素的信息传播长度就是长期依赖性；相当于元素记忆与其相距较远元素的能力，而这种能力体现在长度上；

长期依赖性是语言模型的重要问题之一，RNNs网络通过循环机制让信息进行流通，使后面的元素去记忆前面的元素，但是由于梯度消失和梯度爆炸等问题，导致普通的RNN其依赖性不强，改良后的LSTM，通过门控技术和梯度剪枝技术，能够加强RNN依赖性的能力，但是总体来说其效果并不显著；平均每个LSTM的依赖性长度只有200个词；

Transformer在长期依赖性上面表现得很好，但是其收到了固定长度的限制；在长序列数据中，训练方式只能将长序列分割为多段短序列，然后使用padding转为固定长度，然后每一次丢入固定长度的序列进行训练；如图所示：

很显然这样有一个缺点，就是分割后的多段短序列之间毫无联系，这一问题被称作为context fragmentation，即上下文碎片化；而Transformer-XL就是为了解决这一问题所提出来的模型；

该模型主要的技术贡献在于：在注意力机制中引入了递归的概念，并提出了相对位置编码

2.2 Transformer-XL的架构

从下图中我们可以发现，Transformer-XL利用了循环的机制，在段与段之间建立了联系；相对于普
通的Transformer，从结构上很容易的看出该模型可以优化上下文碎片化这一问题；

与此同时，由于段与段之间进行了连接，模型的长期依赖性就更加强大，并不局限于固定长度的序列；不过在这里有一个问题，就是说普通的Transformer采用的绝对位置编码，段与段之间的位置编码一致，编码出现了重复，这样就会导致信息的传播出现异常，论文中使用相对位置编码解决这个问题；

在训练过程中，普通的Transformer模型每一步只在最后一个位置做一个预测，然后采用right shift，丢掉第一个位置的信息，加入上一次预测的信息，从头开始处理，预测下一个最后一个位置，如此重复；而Transformer-XL可以利用之前计算的信息，每一次预测时只需要信息就可以了，这样处理可以大幅度加快计算速度；如图所示：

这个图有一点儿问题，其尾巴应该是固定的，不是移动的；

2.3 相对位置编码

在运用递归机制之前，首要解决的问题是分段语料位置上的一致性，如果像普通的Transformer使用三角函数绝对编码，其流程如下：
$$\begin{array}{rl}{h}_t& =f(h_{t-1},E_{s_s}+U_{1:L})\\ & \\ h_{t+1}& =f(h_t,E_{s_{s+1}}+U_{1:L})\end{array}$$

从中可以发现每一段语料其位置编码都是一样的，模型无法辨别段与段之间的不同，不知道哪一段在前面，哪一段在后面，造成混淆；

为了深入理解位置编码在注意力层中的作用，我们用数学公式来推导：

$$\begin{array}{rl}& T_i=E_i+U_i\\ & Q_i=W_q{T}_i\\ & K_i=W_k{T}_i\end{array}$$

其中$E_i$是第$i$个的元素的token embedding，$U_i$是第$i$个元素的position embedding，$Q$和$K$是注意力机制的query和key；

计算得到：

$$\begin{array}{rl}{Q}_i^T{K}_j& =(W_q{T}_i{)}^T(W_k{T}_j)\\ & =(W_q(E_i+U_i){)}^T(W_k(E_j+U_j))\\ & =(W_q{E}_i+W_q{U}_i{)}^T(W_k{E}_j+W_k{U}_j)\\ & =\underset{a}{E_i^T{W}_q^T{W}_k{E}_j}+\underset{b}{E_i^T{W}_q^T{W}_k{U}_j}+\underset{c}{U_i^T{W}_q^T{W}_k{E}_j}+\underset{d}{U_i^T{W}_q^T{W}_k{U}_j}\end{array}$$

这里涉及到position embedding的只有三个，分别是b,c,d；其中表达位置信息的是$b$和$c$，表达交错信息的是$d$，为了表达相对的思想，我们需要把所有的$U_i$和$U_j$转化为相对编码；

首先我们要定参考系，而attention机制是围绕query来展开的，所以我们把第$i$个位置的query作为参照物；先把所有$U_j$转为相对编码形式，论文用$R_{i-j}$代替；

$$\begin{array}{rl}{Q}_i^T{K}_j& =\underset{a}{E_i^T{W}_q^T{W}_k{E}_j}+\underset{b}{E_i^T{W}_q^T{W}_k{U}_j}+\underset{c}{U_i^T{W}_q^T{W}_k{E}_j}+\underset{d}{U_i^T{W}_q^T{W}_k{U}_j}\\ & =\underset{a}{E_i^T{W}_q^T{W}_k{E}_j}+\underset{b}{E_i^T{W}_q^T{W}_k{R}_{i-j}}+\underset{c}{U_i^T{W}_q^T{W}_k{E}_j}+\underset{d}{U_i^T{W}_q^T{W}_k{R}_{i-j}}\end{array}$$

接下来要处理的是$U_i$，由于$b,d$已经具有了相对性，而$d$中还有一个绝对位置编码$U_i$，这个信息已经没有实质性作用，这里我们可以把$U_i^T{W}_q^T$换成一个可训练参数$u$；接下来还有$c$中有一个绝对位置编码$U_i$；这里我们也把其换成一个可训练参数$v$；在这种情况下，由于所有查询位置的查询向量都是相同的，因此这表明无论查询位置如何，对不同单词的注意偏差都应该保持不变。

与此同时，$W_k{E}_j$和$W_k{R}_{i-j}$由于其表达的意义不一样，为了更高的自由度，训练更深层次的模型；我们利用$W_{kr}$把$W_k{R}_{i-j}$表示为$W_{kr}{R}_{i-j}$；

最后得到的结果如下：

$$\begin{array}{rl}{Q}_i^T{K}_j& =\underset{a}{E_i^T{W}_q^T{W}_k{E}_j}+\underset{b}{E_i^T{W}_q^T{W}_{kr}{R}_{i-j}}+\underset{c}{u^T{W}_k{E}_j}+\underset{d}{v^T{W}_{kr}{R}_{i-j}}\end{array}$$

在新的公式下，$abcd$都有一个直观的含义：

• $a$ 表示$i$位置和$j$位置内容的交互；
• $b$ 表示$i$位置内容的位置偏差；
• $c$ 表示全局内容的交互；
• $d$ 表示相对位置的偏差；

2.4 分段递归模型

原文的计算公式如下，这里只展示了一个$n：

其中$SG(\cdot )$表示stop gradient，相当于停止训练这一部分，看作常数，$[h_u\circ h_v]$表示$h_u$和$h_v$沿着长度的维度进行拼接；即[length + memory_length, d_model]，这里可能考虑到$GPU$显存的限制会对memory_length进行取舍，在计算$Q$的过程中利用$Q=Q[-length:]$转为原来的维度；而$K$和$V$的维度维持$[length+memor{y}_{l}ength]$不变；这样输出的维度仍然是$[length,d_{m}odel]$；

这样模型就完毕了！

三、过程实现

直接把作者的搬过来了，有时间再写一个tensorflow 2.0 版本的

四、整体总结

Transformer-XL利用recurrent和relative position embedding解决了context fragment的问题，同时大大加快了inference的速度；