Transformers Assemble（PART I ）

寫在前面

經過之前一段時間的 NLP Big Bang，現在相對比較平靜了，Transformer 派已經佔據了絕對的主導地位，在各類應用中表現出色。看標題大家也可以猜個差不多，整理了一系列自《Attention is all you need》之後的對 Vanilla Transformer 的改進論文，和大家一起梳理整個發展過程。這篇是第一趴，都來自ICLR。OK，來看看今天的 Transformers：

• Bi-BloSAN from UTS，ICLR2018
• Universal Transformers from UVA&Google，ICLR2019
• Reformer from Google，ICLR2020

歡迎關注公衆號

BI-DIRECTIONAL BLOCK SELF-ATTENTION FOR FASTAND MEMORY-EFFICIENT SEQUENCE MODELING

這篇論文首先分析了目前幾大類特徵抽取器CNN、RNN、Attention的優缺點，針對其不足提出了一種對self-attention的改進，雙向分塊自注意力機制（bidirectional block self-attention (Bi-BloSA)）。

Masked Block Self-Attention

其最主要的組件是掩碼分塊自注意力機制（masked block self-attention (mBloSA)），基本思想是將序列分成幾個等長的block（必要時進行padding），對每個單獨的block內應用self-attention（intra-block SAN）捕獲局部特徵，然後對所有block輸出再應用self-attention（inter-block SAN）捕獲全局特徵。這樣，每個attention只需要處理較短的序列，與傳統的attention相比可以節省大量的內存。最後通過Context fusion模塊將塊內SAN、塊間SAN與原始輸入結合生成最終上下文表示。整體框架如下圖所示：

Masked Self-Attention

在傳統的self-attention中時序信息是很難編碼進去的，通過對某些位置進行掩碼可以實現位置信息的融合。
$f\left(x_{i}, x_{j}\right)=c \cdot \tanh \left(\left[W^{(1)} x_{i}+W^{(2)} x_{j}+b^{(1)}\right] / c\right)+M_{i j} 1$
而建模雙向信息，只需要將上式最後一項修改即可得到前向和後向SAN
$M_{i j}^{f w}=\left\{\begin{array}{ll} {0,} & {i<j} \\ {-\infty,} & {\text { otherwise }} \end{array}\right.$
$M_{i j}^{b w}=\left\{\begin{array}{ll} {0,} & {i>j} \\ {-\infty,} & {\text { otherwise }} \end{array}\right.$

Bi-directional Block Self-Aattention Network

最終將上述模塊整合後得到的網絡框架如下圖。

實驗工作也非常充分，在九大公開數據集的多類NP任務中取得SOTA，同時模型的計算和內存更高效。

Reference

• Code Here

---

---

UNIVERSAL TRANSFORMERS

上一篇論文針對的是Transformer內存佔用大、對長序列輸入不友好的缺陷，除此之外，這篇論文指出其還存在着以下幾個問題

• 丟棄了RNN的歸納偏置（Inductive Bias），而這對一些任務至關重要（EMNLP 2018的一篇文章對比了Transformer和LSTM得出該結論）。Inductive Bias也可以理解爲Prior Knowledge，可以用來預測從未遇到的輸入的輸出，即模型泛化的能力。（refer：Inductive Bias in Recurrent Neural Networks）。
• 由於固定了Transformer的層數，使得模型計算低效。對於一個輸入，有些部分是比較好理解的，有些部分是比較含糊不清的，如果採取相同的計算會導致無效計算。
• 原始Transformer是非圖靈完備的。（這塊沒仔細看下面不會涉及）

針對以上存在的不足，論文提出了一種Transformer的改進模型：Universal Transformer（UT），如下圖所示。乍一看與Transformer非常相似（不同的地方在下圖用紅框框起來會在下文介紹），下面我們來仔細看看UT是如何解決上述缺陷的。

Parallel-in-Time Recurrence Mechanism

原始的Transformer堆疊層數是固定的（6層或12層），爲了結合RNN的優勢，對每一個token設置了一個Transition Function（參數共享），並且可以一次一次循環迭代計算（Transformer是對attention層後直接輸入一個Dense Layer）。如下圖，橫座標是輸入序列的token，縱座標是循環迭代次數Depth。每一次迭代是在整個序列上進行的，不同於RNN需要從左往右依次計算，大大提高了計算效率。當迭代次數Depth人爲固定時，就變成了原始Transformer。

假設輸入爲長度爲$m$的序列，第$t$次循環後有：
$H^{t}=\text{LAYER NORM }\left(A^{t}+\text { TRANSITION }\left(A^{t}\right)\right)$$A^{t}=\text { LAYERNORM }\left(\left(H^{t-1}+P^{t}\right)+\text { MULTIHEADSELFATTENTION }\left(H^{t-1}+P^{t}\right)\right)$其中Transition Function可以是深度可分離卷積也可以是全連接。

Coodinate Embeddings

在上面公式中有一個$P$，指的是兩種embedding，（position，time）：

1. position embedding和原始Transformer的一樣；
2. timestep embedding，其實本質上和position embedding一樣，只不過這裏是關注於循環步數$t$，因爲不同的循環次數學習的是不同層面的知識，這一信息可以提供給模型學習了幾次，什麼時候該停止；
   $\begin{aligned} P_{i, 2 j}^{t} &=\sin \left(i / 10000^{2 j / d}\right)+\sin \left(t / 10000^{2 j / d}\right) \\ P_{i, 2 j+1}^{t} &=\cos \left(i / 10000^{2 j / d}\right)+\cos \left(t / 100000^{2 j / d}\right) \end{aligned}$

Dynamic Halting

正如前面所說的，輸入中的某些部分需要更多次的計算模型纔可以理解。例如最常見的I went to the bank to deposit money裏bank一詞就需要更多次的運算。在UT中引入了Adaptive Computation Time (ACT)來實現，ACT可以基於模型在每步循環中預測的標量halting probability，動態調製處理每個輸入token所需的計算步驟數（即響應時間ponder time）。UT將動態ACT暫停機制分別應用於每個位置， 一旦特定的循環塊停止，它將其狀態複製到下一個步驟，直到所有塊都停止，或者直到達到最大步數爲止， 編碼器的最終輸出是以此方式產生的最後一層表示。Dynamic Halting策略流程可以概括爲：

1. 在第$t$步循環時，有第$t$步的halting probability、第$t-1$步的token state以及提前設置的超參停止閾值；
2. 利用UT計算當前token state；
3. 利用全連接層和sigmoid激活函數計算ponder time，ponder time表示模型對每個token還需要多少次運算的預測；
4. 判斷每個token當前是否需要停止，規則爲：
   當 (halting probability + pondering) > halt threshold 時，停止；
   當 (halting probability + pondering) ≤ halt threshold 時，繼續運算；
5. 對於那些仍然在運算的token，更新停止概率：halting probability += pondering；
6. 當所有位置都停止或者模型達到循環最大步數時，結束。

整體流程如下面動圖所示。

Reference

• Code Here
• Moving Beyond Translation with the Universal Transformer

---

---

REFORMER: THE EFFICIENT TRANSFORMER

傳統的Transformer仍然存在着一些問題，比如內存佔用大、計算複雜度大、無法較好處理長文本等。針對以上不足，本文作者提出了一種更爲高效的Transformer，該方法的主要目標是能夠處理更長的上下文窗口，同時減少計算量和提高內存效率，具體做法如下:

• 💥局部敏感哈希注意力替代傳統注意力，將複雜度從$O(L^{2})$降低爲$O(L)$；
• 使用可逆層（reversible layers），在訓練過程中只存儲單層激活值，而不是N次；
• 將FF層中的激活進行切分並分塊處理

關於注意力機制

Scaled Dot-Product Attention

Transformer中最關鍵的部分就是其注意力機制，原始自注意力機制通過三個不同的映射矩陣將輸入表示爲Q、K、V三部分，然後計算下式：
$\text { Attention }(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V$詳細解釋可以參考illustrated-transformer、illustrated-self-attention。

Memory-Efficient Attention

但是根據上式計算會發現模型佔用內存非常大，主要問題在於 $QK^{T}$，假定Q、K、V的形狀爲$[batch\_size, length, d_{model}]$，那麼$QK^{T}$項的維度爲$[batch\_size, length, length]$。舉個栗子，對於一個長度爲64K的序列，即使batch_size爲1，這也會是一個$64K*64K$的矩陣，並且序列越長，佔用的內存會越大，這無疑會極大地限制Transformer的性能。但是注意，我們並不需要計算整個Q，而可以選擇對Q中的每一個query $q_{i}$單獨計算即可，然後在需要BP的時候再算一次（emmmm感覺像是工程實現上的trick）。雖然這樣做不那麼computation-efficient，但是比較memory-efficient，可以處理更長的序列。

Locality-Sensitive Hashing Attention

LSH attention使用的是共享Q、K策略，這樣可以有效降低內存佔用，並且下圖實驗顯示共享策略並不會比原始策略效果差。

注意到最終$QK^{T}$是會經過一個softmax函數，而softmax結果的關鍵在於那些比較大的輸入，因此對於每一個query $q_{i}$，我們只需要關注K中與其最接近的元素即可。那麼問題是如何找到最接近的那些呢？這裏就需要用到局部敏感哈希，簡單來說就是使用哈希函數將鄰近的向量以高概率映射到同一個哈希值，而距離遠的向量高概率映射不到同一個哈希值。更具體地，文中使用的LSH哈希策略如下圖：

最終LSH Attention的整體流程如下所示，①局部敏感哈希分桶，不同的顏色代表不同的桶，相同的顏色代表相近的向量；②序列重排，將具有相同顏色的向量放在一起；③切分序列，方便並行化計算；④在分出來的更小的區塊中應用attention，爲了防止溢出向量，會考慮前一個chunk中的元素進行attention

關於內存

雖然通過LSH策略降低了attention的複雜度，但是由於模型通常包含好多層，導致在訓練過程中保存的中間值數量也非常巨大。爲此，Reformer使用了可逆Transformer，不在內存中保存所有的值，而是在反向傳播時按需計算得出，思想借鑑於RevNet的可逆層。

（a）傳統殘差網絡：上一層的激活用於更新下一層的輸入。$y=x+F(x)$
（b）可逆殘差網絡：維護兩組激活值，每層只更新其中的一組。$y_{2}=x_{2}+F(x_{1})$ ;$z_{1}=x_{1}+G(y_{2})$
（c）反向：只需要減去激活值就可以恢復任意中間值

明白了上述思想，在Reformer中就是將Attention層和FeedForward層帶入F和G即可。

reference

• Code Here
• Reformer應用
• Reformer: The Efficient Transformer
• Reddit上關於Reformer的討論

---

---

本來還打算再加上Transformer-XL的，啊太長了就算了吧…
最後大家一定注意防護！！
以上