Transformer XL原理介紹

1. 引言

在自然語言處理中，當前深度學習主流的結構是RNN和Transformer，因爲這兩種結構能夠比較好地捕捉文本的上下文信息，但是，我們知道RNN主要的問題是梯度消失和梯度爆炸的問題，而且其捕捉上下文的長度沒有Transformer那麼強大，而Transformer雖然能力比較強，但是在預測時會受到訓練時所設定的最大長度限制，因此，本文將介紹一個Transformer模型的變體，即Transformer XL(extra long)，該模型是CMU和谷歌大腦在2019年提出來的，通過引入了遞歸的機制和相對位置編碼，解決了Transformer長度限制的問題，作者在實驗中發現，Transformer XL能捕捉到的文本長度比RNN長80%。

• 論文地址：《Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context》

2. Transformer XL原理介紹

2.1 Vanilla Transformer

在介紹Transformer模型之前，先介紹另一個模型，該模型是2018年由Al-Rfou等人提出來的，作者稱之爲Vanilla Transformer，該模型其實本質上還是Transformer，只是將原來的句子進行切片，切成一個一個固定長度的短句子進行訓練，如下圖中的(a)所示，然後在預測時，採用訓練時所設定的長度窗口，在句子上進行平移，從而不斷獲得句子的預測，如下圖中的(b)所示，雖然這樣操作可以解決句子的長度限制問題，但是存在兩個主要的缺點：①每次平移後，模型都是從頭開始計算，並沒有利用前面計算得到的上文信息，因此，會出現一種“上下文斷裂”的情況；②在預測時，當句子的長度比較長時，計算速度會非常慢，因爲每次窗口只移動一個時間步，這樣會導致很多信息重複計算。作者在實驗中發現，採用Transformer XL模型，速度是Vanilla Transformer的1800多倍。

2.2 Segment-Level Recurrence

爲了克服Vanilla Transformer模型的缺點，Transformer XL引入了一種遞歸的機制，即通用將句子按照固定的長度$L$分成若干個子句，然後在訓練和預測時，依次將每個子句傳入Transformer模型，並且將每個子句在Transformer中各層的輸出傳遞給下一個子句，如下圖(a)中綠色線，在下一個子句每一層的計算中，將上一個子句對應上一層的輸出與當前子句對應層的輸入進行拼接，這樣在每一個子句的預測時，就可以考慮到前面各個子句的信息。

記兩個連續的子句爲$\mathbf{s}_{\tau}=\left[x_{\tau, 1}, \cdots, x_{\tau, L}\right]$，$\mathbf{s}_{\tau+1}=\left[x_{\tau+1,1}, \cdots, x_{\tau+1, L}\right]$，記第$\tau$個子句第$n$層的輸出爲$\mathbf{h}_{\tau}^{n} \in \mathbb{R}^{L \times d}$，其中$d$表示隱藏層的維度，則第$\tau+1$個子句$\mathbf{s}_{\tau+1}$第$n$層的輸出計算如下：

$\begin{array}{l}{\widetilde{\mathbf{h}}_{\tau+1}^{n-1}=\left[\mathrm{SG}\left(\mathbf{h}_{\tau}^{n-1}\right) \circ \mathbf{h}_{\tau+1}^{n-1}\right]} \\ {\mathbf{q}_{\tau+1}^{n}, \mathbf{k}_{\tau+1}^{n}, \mathbf{v}_{\tau+1}^{n}=\mathbf{h}_{\tau+1}^{n-1} \mathbf{W}_{q}^{\top}, \widetilde{\mathbf{h}}_{\tau+1}^{n-1} \mathbf{W}_{k}^{\top}, \widetilde{\mathbf{h}}_{\tau+1}^{n-1} \mathbf{W}_{v}^{\top}} \\ {\mathbf{h}_{\tau+1}^{n}=\text { Transformer-Layer }\left(\mathbf{q}_{\tau+1}^{n}, \mathbf{k}_{\tau+1}^{n}, \mathbf{v}_{\tau+1}^{n}\right)}\end{array}$

其中，$\mathrm{SG}(\cdot)$表示stop-gradient，$\left[\mathbf{h}_{u} \circ \mathbf{h}_{v}\right]$表示將兩個序列進行拼接，$\mathbf{W}$表示模型的參數。這裏主要的不同點在於計算K和V時，使用的是$\widetilde{\mathbf{h}}_{\tau+1}^{n-1}$，其中，包含了來自上一個子句的信息$\mathbf{h}_{\tau}^{n-1}$。由於每個子句都會遞歸地使用上一個子句的信息，因此，隨着層數的加深，每個子句中每個時間步可以利用到的前面的序列信息長度應該是$O(N \times L)$，如下圖中(b)所示。在實際訓練中和預測中，對於每個子句，還可以考慮前面多幾個子句的信息，而不是僅使用上一個子句。作者在實驗中通過一個參數$M$來表示考慮的上文的長度信息，在訓練時，令$M=L$，即只考慮一個子句信息，但是在預測時，可以將其擴增幾個倍數。

Transformer XL除了可以捕捉更長的文本信息外，也提高了模型預測的速度，因爲我們知道Transformer XL在預測時每次是平移一個子句的長度，而不是像Vanilla Transformer一樣每次只平移一個時間步。

2.3 Relative Positional Encodings

雖然Transformer XL通過切分子句，並引入遞歸的思想解決了句子長度的限制問題，但是還忽略了一個問題，即positional encoding。在Transformer中，通過一個positional encoding來表徵句子的時序信息，記爲$\mathbf{U} \in \mathbb{R}^{L_{\max } \times d}$，其中，$\mathbf{U}_{i}$表示第$i$個位置的絕對位置信息，$L_{max}$表示句子的最大長度，我們知道，Transformer中，會將詞彙embedding信息與位置信息進行元素相加，然後作爲整個模型的原始輸入，那麼，如果Transformer XL也採用這樣的方式進行計算的話，其表達式如下：
$\begin{aligned} \mathbf{h}_{\tau+1} &=f\left(\mathbf{h}_{\tau}, \mathbf{E}_{\mathbf{s}_{\tau+1}}+\mathbf{U}_{1 : L}\right) \\ \mathbf{h}_{\tau} &=f\left(\mathbf{h}_{\tau-1}, \mathbf{E}_{\mathbf{s}_{\tau}}+\mathbf{U}_{1 : L}\right) \end{aligned}$

其中，$\mathbf{E}_{\mathbf{s}_{\tau}} \in \mathbb{R}^{L \times d}$表示子句$\mathbf{s}_{\tau}$的embedding信息，可以看到，兩個子句$\mathbf{s}_{\tau}, \mathbf{s}_{\tau+1}$用的都是同一個位置信息，這會導致模型沒法區分兩個子句的順序信息。

爲了克服這個問題，作者引入了一種新的Positional Encoding方法，即Relative Positional Encodings。我們知道，Transformer在$q_{i}$和$k_{j}$的注意力時，根據矩陣的分配律，可以得到如下：
$\begin{aligned} \mathbf{A}_{i, j}^{\mathrm{abs}} &=\underbrace{\mathbf{E}_{x_{i}}^{\top} \mathbf{W}_{q}^{\top} \mathbf{W}_{k} \mathbf{E}_{x_{j}}}_{(a)}+\underbrace{\mathbf{E}_{x_{i}}^{\top} \mathbf{W}_{q}^{\top} \mathbf{W}_{k} \mathbf{U}_{j}}_{(b)} \\ &+\underbrace{\mathbf{U}_{i}^{\top} \mathbf{W}_{q}^{\top} \mathbf{W}_{k} \mathbf{E}_{x_{j}}}_{(c)}+\underbrace{\mathbf{U}_{i}^{\top} \mathbf{W}_{q}^{\top} \mathbf{W}_{k} \mathbf{U}_{j}}_{(d)} \end{aligned}$

作者認爲，其實在每次計算時，只要讓模型可以知道句子的相對位置信息即可，因此，基於這個思想，作者將上式中(b)和(d)項中的$\mathbf{U}_{j}$替換爲$\mathbf{R}_{i-j}$，用來表徵相對位置信息，然後將©中的$\mathbf{U}_{i}^{\top} \mathbf{W}_{q}^{\top}$採用一個可訓練的$u \in \mathbb{R}^{d}$向量替代，對於(d)中的$\mathbf{U}_{i}^{\top} \mathbf{W}_{q}^{\top}$，同樣採用一個可訓練的$v \in \mathbb{R}^{d}$替代，最後，對於$\mathbf{W}_{k}$，採用兩個矩陣進行替代$\mathbf{W}_{k, E}$、$\mathbf{W}_{k, R}$。最終修正後的表達式如下：
$\begin{aligned} \mathbf{A}_{i, j}^{\mathrm{rel}} &=\underbrace{\mathbf{E}_{x_{i}}^{\top} \mathbf{W}_{q}^{\top} \mathbf{W}_{k, E} \mathbf{E}_{x_{j}}}_{(a)}+\underbrace{\mathbf{E}_{x_{i}}^{\top} \mathbf{W}_{q}^{\top} \mathbf{W}_{k, R} \mathbf{R}_{i-j}}_{(b)} \\ &+\underbrace{u^{\top} \mathbf{W}_{k, E} \mathbf{E}_{x_{j}}}_{(c)}+\underbrace{v^{\top} \mathbf{W}_{k, R} \mathbf{R}_{i-j}}_{(d)} \end{aligned}$

採用這種相對位置編碼後，假設對於一個$N$層的Transformer，不失一般性，在只有一個multi-heads的情況下，Transformer XL每一層的計算如下：

其中，$\mathbf{h}_{\tau}^{0} :=\mathbf{E}_{\mathbf{s}_{\tau}}$表示詞彙的embedding。雖然採用相對位置編碼使得模型可以區分不同的子句相對位置，但是，我們可以注意到，在計算multi-heads attention時，此時$\mathbf{W}_{k, R}^{n} \mathbf{R}_{i-j}$的計算時間是二次方的，嚴重影響了模型的計算速度，因此，爲了克服這個問題，作者通過推導將其計算時間轉化爲線性的。

我們知道，對於配對$(i, j)$，$i-j$的取值範圍爲$[0,M+L-1]$，因此，可以將$\mathbf{W}_{k, R} \mathbf{R}_{i-j}$的計算轉化爲如下：

此時，對於每一個配對的$(i, j)$，$i-j$，都可以通過位置索引直接從矩陣$Q$中獲取到，爲了便於表示，我們記$\mathbf{Q}_{k}=\mathbf{W}_{k, R} \mathbf{R}_{M+L-1-k}$，此時，對於$\mathbf{q}_{\tau, i}^{n} \top \mathbf{W}_{k, R}^{n} \mathbf{R}_{i-j}$的計算，可以表示如下：

因此，可以定義：

可以發現，對於$\mathbf{B}$中的每一行，其實就是$\widetilde{\mathbf{B}}$中每一行分別向左移動$L、L-1、L-2、...$個單位，因此，通過這樣的方式可以把計算時間轉化爲線性的。

同理，對於$v^{\top} \mathbf{W}_{k, R}^{n} \mathbf{R}_{i-j}$的計算，採用同樣的方法可以定義矩陣：

以及：

此時，對於$\mathbf{D}$中的每一行，其實就是$\widetilde{\mathbf{d}}$向左移動一定單位後的結果。

3. 總結

以上就是Transformer XL的介紹，感覺還是有一定複雜度的。最後總結一下：

• Transformer XL由於加入遞歸機制和相對位置編碼，因此，解決了文本長度的限制問題以及Vanllia Transformer中上下文信息斷裂、預測速度慢的問題。
• 由於採用相對位置編碼，因此，計算速度相比Transformer要慢一點。