[論文] Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context?arxiv.org
Motivation
- Transformer在預訓練階段,設置了固定序列長度max_len的上下文,finetuning階段,模型不能獲取大于max_len的上下文依賴;
- Transformer在長文本編碼過程中,可采用按照句子邊界拆分和按照max_len截斷的方式,進行片段的編碼,論文指出第二種不考慮句子邊界的編碼方式效果更好,但仍然存在上下文碎片的問題context fragmentation。
How to solve
- Segment-level recurrence mechanism - 使模型獲取更長上下文
在單向編碼長文本時,由于transformer設置了max_len,從而在訓練時,第i個token只能attention到前i個token,且當transformer以max_len為窗口滑動時,每前進一步,絕對位置編碼需要跟著前進一步,使得每一步都需要重新計算max_len中每個token的隱層表示,如下圖。
Figure 1: vanilla model with a segment length 4.
Transformer-XL通過設置memory-span使得當前max_len窗口中的每個token都能attention到前max_len個token,因此Transformer-XL在每前進一步時,只用計算當前位置的token的隱層表示,同時在更新梯度時,只更新當前窗口內的梯度(防止梯度bp的距離太遠),從而實現了輸出隱層表示的更長上下文關聯,和高效的編碼速度。
Figure 2: Transformer-XL model with a segment length 4.
- Relative Positional Encodings - 使模型在前向過程中更快
Segment-level recurrence mechanism機制中提到的,max_len窗口中的每個token都能attention前max_len個token,其中可能會有一些token在上一個seg,這樣就存在位置編碼不連續,或者相同token在當前seg和前一個seg具有相同的attention值的問題。因此為了實現transformer-XL訓練和長文本編碼運用之間的等效表示,將絕對位置編碼替換為以當前token為基準的相對位置編碼Relative positional encodings。
其中
分別表示token emb, absolute pos emb, relative pos emb, proj matrix,對于每個編碼的token相對位置編碼都為0,因此
絕對位置編碼在輸入transformer之前就和token emb求和,相對位置編碼需要在計算attention score加入和計算。在Transformer-XL的tensorflow代碼是如何實現呢?
Relative positional emb 代碼解析
- 在解析代碼前,先用圖示展示relative pos emb的實現過程(無memory版本) rel_shift(*)
- 輸入token emb和反向的absolute pos emb
- 得到attention score矩陣后,在token emb維pad,產生1位錯位;
- 截取位置編碼對齊后的矩陣。
- 按順序截取token emb個數個分數組成行,對角全是pad
- 在tf中,rel_multihead_attn函數生成相對位置的attention score
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')
# w : token emb
# r : 反向的絕對位置emb
# r_w_bias :公式中的u
# r_r_bias : 公式中的v
# attn_mask : attention mask矩陣
# mems : memory
def rel_shift(x):x_size = tf.shape(x)x = tf.pad(x, [[0, 0], [1, 0], [0, 0], [0, 0]]) #第二維padding [qlen,klen,bsz,nhead]x = tf.reshape(x, [x_size[1] + 1, x_size[0], x_size[2], x_size[3]]) #reshape產生偏移x = tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1]) #截取attention score矩陣x = tf.reshape(x, x_size)return x
Segment-level recurrence mechanism - Memory 代碼解析
- 同樣在解析代碼前,先用圖示展示Memory的實現過程
- 當前 為query,和 內積;
- 按照之前講解的方式得到relative pos和 的內積結果;
- 得到attention score后,通過 得到attention score矩陣 ;
- 在tf中,_cache_mem(*)函數返回上一個
def _cache_mem(curr_out, prev_mem, mem_len=None):if mem_len is None or prev_mem is None:new_mem = curr_outelif mem_len == 0:return prev_memelse:new_mem = tf.concat([prev_mem, curr_out], 0)[- mem_len:]return tf.stop_gradient(new_mem)
def _create_mask(qlen, mlen, same_length=False): #same_length : 每個token是否采用相同長度的attn length# 代碼中train階段為False 測試時是Trueattn_mask = tf.ones([qlen, qlen]) # 1: mask 0: non-maskmask_u = tf.matrix_band_part(attn_mask, 0, -1) #上三角 = 1mask_dia = tf.matrix_band_part(attn_mask, 0, 0) #對角 = 1attn_mask_pad = tf.zeros([qlen, mlen]) # memory的maskret = tf.concat([attn_mask_pad, mask_u - mask_dia], 1) #如果使token相同的attn_len,設置下三角maskif same_length:mask_l = tf.matrix_band_part(attn_mask, -1, 0)ret = tf.concat([ret[:, :qlen] + mask_l - mask_dia, ret[:, qlen:]], 1)return ret
- relative pos encoding完整代碼
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 ww_heads = tf.layers.dense(cat, 3 * n_head * d_head, use_bias=False,kernel_initializer=kernel_initializer, name='qkv')# word線性映射r_head_k = tf.layers.dense(r, n_head * d_head, use_bias=False,kernel_initializer=kernel_initializer, name='r')# pos線性映射w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, -1)w_head_q = w_head_q[-qlen:] #將memory從query中剔除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_biasrr_head_q = w_head_q + r_r_biasAC = 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) * scaleattn_mask_t = attn_mask[:, :, None, None]attn_score = attn_score * (1 - attn_mask_t) - 1e30 * attn_mask_tattn_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
創作挑戰賽新人創作獎勵來咯,堅持創作打卡瓜分現金大獎
總結
以上是生活随笔為你收集整理的transformer机制讲解_【核心代码解读】Transformer-XL的全部內容,希望文章能夠幫你解決所遇到的問題。
如果覺得生活随笔網站內容還不錯,歡迎將生活随笔推薦給好友。