[Transformer-XL]论文实现：Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context
文章目录
-
- 一、完整代码
- 二、论文解读
-
- 2.1 Transformer-XL的介绍
- 2.2 Transformer-XL的架构
- 2.3 相对位置编码
- 2.4 分段递归模型
- 三、过程实现
- 四、整体总结
论文：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-XL]论文实现：Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context_第1张图片](http://img.e-com-net.com/image/info8/3f21aa71a93a4b8fb3c3e71663b96a70.jpg)
与此同时，由于段与段之间进行了连接，模型的长期依赖性就更加强大，并不局限于固定长度的序列；不过在这里有一个问题，就是说普通的Transformer采用的绝对位置编码，段与段之间的位置编码一致，编码出现了重复，这样就会导致信息的传播出现异常，论文中使用相对位置编码解决这个问题；
在训练过程中，普通的Transformer模型每一步只在最后一个位置做一个预测，然后采用right shift，丢掉第一个位置的信息，加入上一次预测的信息，从头开始处理，预测下一个最后一个位置，如此重复；而Transformer-XL可以利用之前计算的信息，每一次预测时只需要信息就可以了，这样处理可以大幅度加快计算速度；如图所示：
![[Transformer-XL]论文实现：Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context_第2张图片](http://img.e-com-net.com/image/info8/20a32c7bb6c74a0c8880ce926ad3a741.gif)
这个图有一点儿问题，其尾巴应该是固定的，不是移动的；
2.3 相对位置编码
在运用递归机制之前，首要解决的问题是分段语料位置上的一致性，如果像普通的Transformer使用三角函数绝对编码，其流程如下：
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 ) \begin{aligned} 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{aligned} htht+1=f(ht−1,Ess+U1:L)=f(ht,Ess+1+U1:L)
从中可以发现每一段语料其位置编码都是一样的，模型无法辨别段与段之间的不同，不知道哪一段在前面，哪一段在后面，造成混淆；
为了深入理解位置编码在注意力层中的作用，我们用数学公式来推导：
T i = E i + U i Q i = W q T i K i = W k T i \begin{aligned} &T_{i} = E_{i} + U_{i} \\ &Q_i = W_qT_{i} \\ &K_i = W_kT_{i} \\ \end{aligned} Ti=Ei+UiQi=WqTiKi=WkTi
其中 E i E_i Ei是第 i i i个的元素的token embedding， U i U_i Ui是第 i i i个元素的position embedding， Q Q Q和 K K K是注意力机制的query和key；
计算得到：
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 ) = E i T W q T W k E j ⏟ a + E i T W q T W k U j ⏟ b + U i T W q T W k E j ⏟ c + U i T W q T W k U j ⏟ d \begin{aligned} Q_i^TK_j &= (W_qT_{i})^T(W_kT_{j}) \\ &=(W_q(E_{i} + U_{i}))^T(W_k(E_{j} + U_{j})) \\ &=(W_qE_{i} + W_qU_{i})^T(W_kE_{j} + W_kU_{j}) \\ &=\underbrace{E_{i}^TW_q^T W_kE_{j}}_a + \underbrace{E_{i}^TW_q^T W_kU_{j}}_b + \underbrace{U_{i}^TW_q^TW_kE_{j}}_c + \underbrace{U_{i}^TW_q^TW_kU_{j}}_d \end{aligned} QiTKj=(WqTi)T(WkTj)=(Wq(Ei+Ui))T(Wk(Ej+Uj))=(WqEi+WqUi)T(WkEj+WkUj)=a EiTWqTWkEj+b EiTWqTWkUj+c UiTWqTWkEj+d UiTWqTWkUj
这里涉及到position embedding的只有三个，分别是b,c,d；其中表达位置信息的是 b b b和 c c c，表达交错信息的是 d d d，为了表达相对的思想，我们需要把所有的 U i U_i Ui和 U j U_j Uj转化为相对编码；
首先我们要定参考系，而attention机制是围绕query来展开的，所以我们把第 i i i个位置的query作为参照物；先把所有 U j U_j Uj转为相对编码形式，论文用 R i − j R_{i-j} Ri−j代替；
Q i T K j = E i T W q T W k E j ⏟ a + E i T W q T W k U j ⏟ b + U i T W q T W k E j ⏟ c + U i T W q T W k U j ⏟ d = E i T W q T W k E j ⏟ a + E i T W q T W k R i − j ⏟ b + U i T W q T W k E j ⏟ c + U i T W q T W k R i − j ⏟ d \begin{aligned} Q_i^TK_j &=\underbrace{E_{i}^TW_q^T W_kE_{j}}_a + \underbrace{E_{i}^TW_q^T W_kU_{j}}_b + \underbrace{U_{i}^TW_q^TW_kE_{j}}_c + \underbrace{U_{i}^TW_q^TW_kU_{j}}_d \\ &=\underbrace{E_{i}^TW_q^T W_kE_{j}}_a + \underbrace{E_{i}^TW_q^T W_kR_{i-j}}_b + \underbrace{U_{i}^TW_q^TW_kE_{j}}_c + \underbrace{U_{i}^TW_q^TW_kR_{i-j}}_d \end{aligned} QiTKj=a EiTWqTWkEj+b EiTWqTWkUj+c UiTWqTWkEj+d UiTWqTWkUj=a EiTWqTWkEj+b EiTWqTWkRi−j+c UiTWqTWkEj+d UiTWqTWkRi−j
接下来要处理的是 U i U_i Ui，由于 b , d b,d b,d已经具有了相对性，而 d d d中还有一个绝对位置编码 U i U_i Ui，这个信息已经没有实质性作用，这里我们可以把 U i T W q T U_{i}^TW_q^T UiTWqT换成一个可训练参数 u u u；接下来还有 c c c中有一个绝对位置编码 U i U_i Ui；这里我们也把其换成一个可训练参数 v v v；在这种情况下，由于所有查询位置的查询向量都是相同的，因此这表明无论查询位置如何，对不同单词的注意偏差都应该保持不变。
与此同时， W k E j W_kE_{j} WkEj和 W k R i − j W_kR_{i-j} WkRi−j由于其表达的意义不一样，为了更高的自由度，训练更深层次的模型；我们利用 W k r W_{kr} Wkr把 W k R i − j W_kR_{i-j} WkRi−j表示为 W k r R i − j W_{kr}R_{i-j} WkrRi−j；
最后得到的结果如下：
Q i T K j = E i T W q T W k E j ⏟ a + E i T W q T W k r R i − j ⏟ b + u T W k E j ⏟ c + v T W k r R i − j ⏟ d \begin{aligned} Q_i^TK_j &=\underbrace{E_{i}^TW_q^T W_kE_{j}}_a + \underbrace{E_{i}^TW_q^T W_{kr}R_{i-j}}_b + \underbrace{u^TW_kE_{j}}_c + \underbrace{v^TW_{kr}R_{i-j}}_d \end{aligned} QiTKj=a EiTWqTWkEj+b EiTWqTWkrRi−j+c uTWkEj+d vTWkrRi−j
在新的公式下， a b c d abcd abcd都有一个直观的含义：
- a a a 表示 i i i位置和 j j j位置内容的交互；
- b b b 表示 i i i位置内容的位置偏差；
- c c c 表示全局内容的交互；
- d d d 表示相对位置的偏差；
2.4 分段递归模型
原文的计算公式如下，这里只展示了一个$n：
![[Transformer-XL]论文实现：Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context_第3张图片](http://img.e-com-net.com/image/info8/9a3611025fb740a886bea3f1c82fee44.jpg)
其中 S G ( ⋅ ) SG(·) SG(⋅)表示stop gradient，相当于停止训练这一部分，看作常数， [ h u ∘ h v ] [h_u\circ h_v] [hu∘hv]表示 h u h_u hu和 h v h_v hv沿着长度的维度进行拼接；即[length + memory_length, d_model]，这里可能考虑到 G P U GPU GPU显存的限制会对memory_length进行取舍，在计算 Q Q Q的过程中利用 Q = Q [ − l e n g t h : ] Q = Q[-length:] Q=Q[−length:]转为原来的维度；而 K K K和 V V V的维度维持 [ l e n g t h + m e m o r y l e n g t h ] [length + memory_length] [length+memorylength]不变；这样输出的维度仍然是 [ l e n g t h , d m o d e l ] [length, d_model] [length,dmodel]；
这样模型就完毕了！
三、过程实现
直接把作者的搬过来了，有时间再写一个tensorflow 2.0 版本的
四、整体总结
Transformer-XL利用recurrent和relative position embedding解决了context fragment的问题，同时大大加快了inference的速度；