A Brief Overview: Vanilla Transformer v.s. Universal Transformer

無所不能的通用計算模型

Preface

2018年7月，Google Brain 基於Attention is all you need[1] 提出了泛用型Transformer，除了在機器閱讀理解和機器翻譯有了新的突破，也彌補Vanilla (傳統) Transformer在算數及基本邏輯任務上表現不好的問題，透過Universal Transformer，我們可以打破模型架構上的限制，讓模型能針對不同任務的類型、難易，有不同的學習策略和學習時間，因此能泛用於所有任務且達到相當不錯的表現，此篇文章將針對Universal Transformer[2]技術細節進行探討，並說明為何此架構能具有通用計算能力。

Encoder-Decoder Review

在介紹Universal Transformer之前，我們先來回顧一下傳統Encoder-Decoder的作法：

在自然語言處理領域中，有許多應用可透過Encoder-Decoder模型來實現，例如機器翻譯(Machine Translation), 文本摘要(Text Summarization), 聊天機器人(Chat-bot)…等等，這些任務都有一個共同的特性，那就是模型的encoding state需輸入一段文字，咀嚼過這些輸入文字後，接著在decoding state輸出對應的文字反饋給使用者。

以中翻英機器翻譯為例，如果encoder輸入「我愛你」則decoder應輸出「I love you 」，若要訓練一個聊天機器人，如下圖所示，則首先將「How are you」按照詞彙排列順序，依序作為encoder的輸入，接著我們基於encoder last hidden state, output作為decoder state的initial state，隨後基於decoder的輸出做為下一個時間點的輸入，直至輸出達到長度的上限或停止符號<EOL>即完成輸出。

可惜上述的做法存在一些問題：

首先Recurrent Neural Network(RNN) based方法，當輸入長度太長時，容易在資訊傳遞的過程有資訊遺失的狀況，例如:

「小明早上和朋友去河堤玩，中午一起去麥當勞吃飯順便討論報告，回家後與爸媽一同逛夜市，他過了開心又充實的一天。」

此例是自然語言處理中常見的指代消解(Coreference Resolution)問題，代名詞「他」等同於「小明」，但由於兩者的距離太遠，而模型每個時間點只能輸入一個字元，導致訓練的過程中，模型沒辦法充分理解「他」與「小明」的關聯。

此外，在深層的網路架構中，反向傳播(Backpropagation)容易有梯度爆炸(Gradient exploding)和梯度消失(Gradient vanishing)的現象，這也是造成傳統Encoder-Decoder表現低落的主要因素。

Encoder-Decoder with Attention Mechanism

為了解決RNN先天上的限制，2015年Yoshua Bengio在JMLR發表了一篇論文 Show, Attend and Tell: Neural Image Caption Generation with Visual Attention[3] 將注意力機制(Attention Mechanism)使用在影像轉文字上，大幅地提升預測的準確度，不再受序列和區域的侷限；注意力機制也因此聲名大噪，各方面基於注意力機制的模型就如同雨後春筍般地冒出頭來。

Encoder-Decoder with Attention Mechanism 架構就如下圖所示：

注意力機制模型中，Decoder 的輸出會與Encoder每一個時間點的輸出進行分數計算(如下公式)，分數計算方法很多，常見的方式為內積(dot product)，另外也可以乘上一個可訓練的權重W (general) 或合併(concatenation)再乘；接著通過softmax層便可得到每一個encoder hidden state’s attention weights其總和為1，可表示為該decoder步驟注重encoder 輸出的權重分布，隨後我們將每個encoder 輸出與 attention weights 作加權總合即得到最終的attention vector，Decoder在每一個時間點都會進行上述的注意力計算，直至結束輸出為止。

上述注意力機制，搭配LSTM、GRU cell 已可大大提升整體表現，不僅如此，也能透過注意力權重，了解機器在Decoder的每個時間點，著重Encoder的那些部分，增加了模型的可解釋性；但如同前面所說，RNN這種需要隨著序列一個一個將字元送進模型的作法，在先天上還是有很大的劣勢，因此在Facebook於2017年中提出 ConvNet Seq2seq[4] 讓序列問題可以併行處理後，Google隨後提出了震撼彈 Transformer 於 Attention is all you need 論文中，讓所有目光又聚焦到了Google身上。

Attention is all you need

Transformer 有別於傳統方法，摒棄了序列任務所需的多重步驟，取而代之，模型可進行一次性輸入所有的文字，並平行化的進行訓練，讓每一次的訓練不再依賴前一個時間點的狀態，不僅加快整體訓練速度，也提升任務表現，在機器翻譯的任務 WMT 2014 English-to-German[5] 取得 BLEU 41.8，在當時已超越前一個state-of-the-art 2個百分點，成為當時最好的機器翻譯模型。

Transformer 雖然厲害但它還是有弱點，因為架構上的侷限，在算數及基本邏輯任務上表現都遠劣於傳統的RNN模型，此外Transformer 中 Encoder & Decoder Block執行的次數仍是人為設定的magic number，這意味著，模型架構無法隨著任務的特性、難易，自動的調整訓練的策略和次數，因此今年七月，Google便發表了Universal Transformer來彌補前者的缺失，成為更通用的encoder-decoder模型，以下便開始進入正題，針對Universal Transformer技術細節進行探討。

Universal Transformer Architecture

上面為Universal Transformer架構圖，先快速瀏覽一下每個部份：

• Universal Transformer 如往常一樣，分為Encoder Block & Decoder Block 兩個部分，每個Block都會遞迴執行T次。
• Universal Transformer 透過正弦函數具週期的特性，作Position embedding將位置資訊加入每一個input symbol，除此之外，Timestep embedding(遞迴次數)也使用相同方法處理。
• Universal Transformer 透過 Multihead Self-Attention 來計算每個字對每個字的注意力權重，其中Decoder的Multihead Self-Attention在訓練階段會增加遮罩，確保目前要預測的position只注意到前面已預測過的字元，避免偷看到後面尚未預測的答案。
• 使用 Layer normalization 和 Dropout 來避免模型訓練過擬合(Overfitting) 和 Internal Covariate Shift(ICS)[6]的情況並增加模型的泛化能力，提升訓練速度。
• 使用Transition function，除了增加非線性變換的次數外，與Vanilla Transformer不同的是，每一個Symbol都有自己獨立的Transition function，這能夠有效率的學習到更深層的語意資訊。

下面將針對上述作法進行詳細解說：

Position & Timestep Embedding

正弦函數Cosine & Sine Function具週期性，因此將每一個symbol加上由正弦函數所建構的Positional encoding藉以添加相對位置資訊，其中偶數維度用sine函數，奇數維度則使用cosine函數。

pos: symbol position, j: dimension, d: input dimension(fixed), t:iteration times

當2j為0時，函數以2 π為周長，若仔細觀察下列公式會發現，2j範圍從0至d，意思是每一個symbol的vector維度是d，每一個維度都會有不同的週期(周長)，從0~d，周長變化呈幾何級數成長，從2 π至10000 · 2π。

另外在不考慮iteration encoding項的情況下，我們可找到任一offset k，使得位置pos+k可透過線性轉換表示成pos如下證明：

P_(pos+k, 2j)

= sin((pos+k)/a)

= sin(pos/a + k/a)

= cos(pos/a)sin(k/a) + sin(pos/a)cos(k/a)

= P_(pos, 2j+1)u + P_(pos, 2j)v

= (P_(pos, 2j+1), P_(pos, 2j)) . (u, v)

同理在不考慮positional encoding的情況下，我們可以找到任一offset k，使得t+k可透過線性轉換表示成t，如下證明：

P_(2j+1)^(t+k)

= cos(t/b + k/b)

= cos(t/b)cos(k/b) + sin(t/b)sin(k/b)

= P_(2j+1)^tu + P_2j^tv

= (P_2j^t , P_(2j+1)^t). (u, v)

最後我們將position & timestep embedding 一併與輸入作加法運算，即成功將位置和遞迴次數資訊加入symbol中；在Attention is all you need 論文中，作者有實驗過將positional embedding function 加入可訓練參數，訓練上述公式heuristic number (e.g. 10000)，最後發現，採用可訓練的參數其結果與上式差不多，便直接採用，且沿用至此篇。

與Vanilla Transformer不同的是，加入了timestep embedding，因為不同的學習次數可學習不同層面的知識，有了timestep embedding，機器就可以知道自己學習了幾輪，應該在什麼時間點停止。

position & timestep embedding 看起來雖然簡單、直觀，但它是這個架構不可或缺的一部分，也是取代RNN方法，同時保留位置資訊的一項重要技巧。

Multi Head Self-Attention

self-attention是Universal Transformer架構的重中之重，摒棄了RNN & CNN，作者標榜只用注意力機制即可讓模型有良好的特徵學習，如上圖所示，其實self-attention與CNN有點像，只是CNN被kernel(window) size給限制，但self-attention可一次看到所有的資訊，並從所有資訊中，找到每個字應該注意的對象，以下例，經過self-attention的注意力權重計算，我們可以發現「its」注意到了「Law」和「application」，一個是它的同位詞，另一個則是所有格後面接的名詞。

來探究一下self-attention的公式，首先將input作linear projection個別乘上W^Q, W^K, W^V，產生 Q, K , V相同維度的向量，接著將Q, K作內積除以輸入維度d^1/2作normalization，增加訓練穩定度和收斂速度，之後通過softmax層，產生m×m 的 attention weight matrix，最後將matrix與V相乘(attend)，即完成self-attention的流程，其中MultiHeadSelfAttention為k個self-attention並行計算，Concat起來後再作最後到線性轉換乘以W^O得到最終結果，我們可以發現H^0經過一連串的運算後，維度仍為m×d。

Q, K, V代表的意含，我們可以把Q(query)視為提出的問題或請求，K(key)基於Q的請求進行計算(內積)，挑選出Q所需資訊對應的索引值，V(value)則為索引值對應的實際內容。

以下Transformer架構為例，假設在inference time，已輸出「我」、 「喜歡」二詞，則接續預測下一個字時，Q就會發出請求，請求內容為『希望下一個輸出是個受詞』，而K就會基於Q的需求進行計算，最後從V中挑選出此階段應該要特別注意的資訊。

Residual Connection

殘差網路(Residual Connection)[7]，其原理如下，就是淺層網路可以直接略過中間的連結，直接與深層網路作skip connection，其好處主要有二：

1. 參數可直接從深層傳遞至淺層，不必再層層通過，除了可加快收斂速度外，亦可讓淺層網路有well-trained的效果。
2. 殘差網路避免了層層的傳輸，因此深層網路架構在反向傳播時可減緩梯度爆炸或梯度消失現象。
3. 作者主要的核心想法是：假設目標是求得輸出H(x)，在某些情況，我們發現直接求得H(x)並非這麼容易，因此我們轉而去求H(X)的殘差F(X) = H(X) - X，這意味著原本必須要求得H(X)，現在我們只需要計算輸入X和輸出H(X)的差距F(X) =H(X) -X，因此事情就變得容易許多。

Dropout

Dropout 是實現regularization的一種方法，於2014年被 Srivastava[8]等人提出，Dropout 在訓練階段隨機的忽略一定數量(自行設定)的神經元，讓神經網路在訓練的階段，只隨機靠部分的神經元進行訓練，此方法讓模型部會過度仰賴某些神經元，可有效緩解過擬合的現象，是目前Deep Learning十分常用的方法。

Layer Normalization

Layer normalization[9]與Batch normalization類似，都是將資料分布轉換為平均為0、標準差為1的高斯分布，Batch normalization是以批次的資料作為轉換的母體，例如Batch size為128，則我們將每128筆資料的平均數、標準差轉換為0和1；而Layer normalization是以神經網路中每層的神經元數量作為轉換目標，例如神經網路中某層有256個神經元，則我們會將這256個神經元的輸入作上述的操作，將該層輸入的分布轉換為高斯分布。

Normalization的好處在於，在神經網路中數據的傳遞，可將所有數值normalize成平均數：0、標準差：1，因此每個神經元輸出的scale相等，具有相同的影響力，解決反向傳播過程中梯度爆炸或梯度消失的問題。

encoder-decoder由於訓練過程中有時間順序的關係，因此無法使用Batch normalization，這時Layer normalization便可基於單一筆訓練資料中各層的神經元，而非基於批次資料來進行歸一化，因此特別適用於encoder-decoder中常見的任務。

Transition Function

Transition function 是有別於Vanilla Transformer一個重要的特點，Transition function可以是一個簡單的Dense Layer，也可以是CNN Layer，差別在於，Vanilla Transformer再通過self-attention後，將所有的輸出一併送進Dense Layer，而Universal Transformer的Transition function則是基於每一個不同的symbol分頭進行。

作個比方，傳統的方法就像把所有的魚、肉、菜丟進一個大鍋爐，而Transition function則根據不同的食材分頭料理，這樣的好處在於，我們可以透過分頭進行，有效率地將每個symbol深層的語意進行充分理解，而非在一開始似懂非懂的階段，就把所有東西丟進全連接層；透過此方法，能提升訓練收斂速度且讓每個symbol有更好的vector representation，因此能提升整體表現。

Conditional computation

每一個symbol在著訓練的過程，都能自由的控制iteration次數，因為有些symbol需要多次的iteration才能充分理解其語意，有些則不需要，先停止的symbol會複製相同的數據到下一個step，直到所有的symbol都停止，或到達了iteration的上限值。

“I went to the bank to deposit money”

以上句為例，句子中的所有詞彙除了bank，都能快速的掌握其語意資訊，但bank基於不同的上下文會有不同的意思，bank可以是河堤也可以是銀行，因此模型訓練時，可以動態地調整每一個symbol block iteration次數，讓模型能透過多次的迴圈中學習到bank是指銀行，而非河堤。

• Difference with Transformer

與Vanilla Transformer不同的地方有二，首先傳統方法迴圈次數是人工設定，而Universal Transformer則是可學習的參數，這意味著前者無法隨著不同任務的特性和難度自由調整學習迴圈次數，而後者可以；接著傳統方法迴圈次數 是一視同仁，所有的symbols都被設定為相同的迴圈次數，而後者則可隨著不同symbol的理解程度動態調整迴圈次數，也因次Universal Transformer變得更全面、更彈性，對於不同的任務採取不同的學習策略，因此作者將之稱為通用計算模型。

Training & Inference Time

Training time

Encoding階段將所有symbol一併送入block內進行self-attention和上述一連串的運算，Decoding階段則依照不同時間點，將待預測symbol的前面所有symbol輸入至decoder block，但在實務上會一併把所有symbol一次性輸入，只是會通過遮罩，確保當前待預測的位置，不會偷看到後面的答案，如下圖所示，我們會把待訓練的資料此方式將上三角的部分屏蔽掉，也因此訓練階段強迫採用teacher forcing，以前面預測的ground truth，作為下一個時間點預測的輸入。

Inference Time

預測階段和一般RNN-based encoder-decoder 雷同，將前面階段的輸出作為下一個時間點的輸入，此作法和Training time採用teacher forcing的行為不同，也因此會有訓練與預測階段行為不一致的情況。

Experiments

bAbI Question-Answering

bAbI QA [10]資料集給定幾個事實描述後，基於這些事實進行回答，分為20種不同的問題類型，作者將每一個事實句當作一個symbol，事實句向量則為該句的詞彙向量加總，從下圖可以看出，不同的事實句理解的時間都不一樣。

從下面實驗數據可以看到，Vanilla Transformer不擅於處理此任務，表現遠落後於其他state-of-the-art model，而Universal Transformer改善了前者的缺點，不僅在precision上有很大的提升更一舉擊敗了其它competitive model成為新的state-of-the-art。

其中bAbI共20種不同類型任務，若單一任務error rate高於5%則記為不及格，adapt. 表示有採用 adaptive computation time。

Subject-Verb Agreement

subject-verb agreement[11] 要判斷句子是否合乎文法，例如判斷單數/複數，動詞該不該加s等等，如下圖，Number of attractors 是指主詞和動詞中間有幾個名詞，越多的attractors也代表越容易受混淆，下面四個例子中，attractor個數分別為0, 1, 1, 3；實驗數據發現，Adaptive Universal Transformer得到最好的結果。

LAMBADA task

LAMBADA[12] 為語言模型任務，目標為給定一個段落，預測問題中挖空的詞彙，其中挖空詞彙可能來自段落、也可能不在段落中；可以發現，Universal Transformer不管在perplexity或accuracy都獲得最好的表現。

Algorithmic Tasks

算術問題如複製相同字串(Copy)、字串複製兩次(Double)、加法(Addition)、字串反過來寫(Reverse)...等等，這類的問題由於Vanilla Transformer架構上的限制，先前的表現，都不如傳統LSTM的作法，而改進後的Universal Transformer則大獲全勝，獲得壓倒性的優勝，除了輸入程式碼輸出計算結果這種較困難的任務外，其他的任務，都可以良好的學習到其中的規則，獲得100%的accuracy。

Conclusion

我們來總結一下這篇paper，Universal Transformer採用了 adaptive computation time (ACT)讓模型可以動態地調整每個symbol的計算次數，藉以獲得最好的向量表達，此外有別於傳統作法，此方法在transition function依照個別symbol分頭進行，這能夠有效率的學習每個symbol的向量表達，並保留了Vanilla Transformer的優點，解決傳統方法梯度爆炸/消失和資訊遺失的問題；實驗發現Universal Transformer不僅讓機器翻譯達到更準確的效果，也補足了algorithmic tasks 成效低落的缺點，讓Universal Transformer能更泛用於各式各樣不同的任務，成為計算通用模型。

References

1. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, 2017.
2. Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, Łukasz Kaiser. Universal Transformer, arXiv:1807.03819v1 [cs.CL] 10 Jul 2018.
3. Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhutdinov, Richard Zemel, and Yoshua Bengio. Show, Attend and Tell: Neural Image Caption Generation with Visual Attention. In International Conference on Machine Learning, 2015.
4. Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, Yann N. Dauphin, Convolutional Sequence to Sequence Learning, arXiv preprint arXiv:1705.03122, 2017.
5. Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. abs/1502.03167, 2015.
6. WMT 2014 English-to-German http://www.statmt.org/wmt14/translation-task.html
7. Szegedy, C.; Ioffe, S.; and Vanhoucke, inception-resnet and the impact of residual connections on learning. arXiv preprint arXiv:1602.07261, 2016.
8. Nitish Srivastava, Geoffrey E Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15(1):1929–1958, 2014.
9. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
10. Jason Weston, Antoine Bordes, Sumit Chopra, Alexander M Rush, Bart van Merriënboer, Armand Joulin, and Tomas Mikolov. Towards ai-complete question answering: A set of prerequisite toy tasks. arXiv preprint arXiv:1502.05698, 2015.
11. Tal Linzen, Emmanuel Dupoux, and Yoav Goldberg. Assessing the ability of lstms to learn syntax-sensitive dependencies. Transactions of the Association of Computational Linguistics, 4(1):521–535, 2016.
12. Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernandez. The lambada dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), volume 1, pages 1525–1534, 2016.