圖解 Transformer

注：原文作者 Jay Alammar，地址 https://jalammar.github.io/illustrated-transformer/

在先前文章中，我們介紹了能大大提高神經機器翻譯的效果 Attention 結構，它目前也普遍存在於現代深度學習模型中。本文將要介紹的 Transformer，則是利用 Attention 來提高這類模型訓練速度的方法。 在特定任務中，Transformers 表現已經優於谷歌的神經機器翻譯。而它最大的優勢則是來自於本身可以並行化。谷歌也推薦使用它作為 Cloud TPU 的教學範例。因此本文將結構這一模型的各個部分，並逐一解釋其工作原理。

Transformer 起初由 Attention is All You Need 這篇文章提出。Tensor2Tensor 庫中則提供了基於 TensorFlow 的實現，而哈佛的 NLP 組給出了 PyTorch 版本的指南，以註釋形式解讀原文。本文我們將高度簡化地逐一介紹模型原理，使對這一主題涉獵不多的讀者也能理解。

觀其大略

我們不妨想把模型看作一個黑盒子，在機器翻譯場景下，我們輸入一個句子，另一端則會輸出翻譯的結果。

檢視模型內部結構，我們發現有一個編碼元件，一個解碼元件，二者之間有連線線。

編碼元件是一批堆疊在一起的編碼器（文章中使用了 6 個編碼器，選擇 6 個沒有什麼玄學原因，換其他數量未嘗不可）。而解碼元件則是由相同數量的解碼器堆疊而成。

所有編碼器都有著相同的結構（然而它們並不共享權重），每個都可以拆分為兩個子結構：

編碼器的輸入首先會進入 self-attention 層，在編碼輸入句子的某一個詞語時，該層會納入句中其他單詞的資訊。下文會詳述這一過程。

self-attention 層的輸出會被餵給一個前饋神經網路，句中不同位置的輸出會各自獨立地被送往相同的前饋網路中。

解碼器也包含了這兩層結構，但在二者之間多了一個 attention 層來幫助其聚焦於輸入句中的相關部分（attention 的作用同其在seq2seq 模型中的類似）。

是時候來點張量了

既然我們已經瞭解模型的主要結構，下面我們來關注下其中涉及的各種向量，張量，看看它們在元件間如何傳遞使得我們能得到想要的模型輸出。

同大多數 NLP 模型一樣，我們先要利用embedding algorithm將輸入的文字轉化為數值向量。

每個詞語都被嵌入成了一個長度為 512 的向量，圖中將用小方塊來表示它們。

詞嵌入過程只在最底層的編碼器中執行一次，而每個編碼器都接受一個由長度為 512 的向量組成的列表作為輸入。對於最底層的編碼器而言，輸入就是句中每個詞嵌入後的向量列表。其他上層編碼器的輸入則是其正下方編碼器的輸出，輸出列表的長度是我們可以設定的一個超引數，一般而言我們會將其設定成訓練集中最長句子的長度。

經過詞嵌入過程後的每一個輸入句，都會流經編碼器的這兩層結構。

現在我們發現了 Transformer 的一個核心特性，輸入句中不同位置的詞語在編碼器中的計算路徑是不同的。在 self-attention 層中這些路徑之間相互關聯，而在前饋網路層中則是獨立的，因此前饋層中的計算過程是可以並行的。

接下來我們將用一個更短的句子為例，看看每個編碼器的子層都做了什麼。

來編碼吧！

正如上文提到的，編碼器接收由向量組成的列表，然後將向量依次輸送至 self-attention 層和前饋神經網路，再將輸出傳至下一個編碼器。

每個位置上的詞語都要經過這樣的 self-attention 處理，然後各自傳向一個同一個前饋神經網路。

Self-Attention 概述

假定我們想翻譯如下句子：

”The animal didn't cross the street because it was too tired”

句中的 “it” 指代什麼呢，是街道還是動物? 對於人類而言這個問題輕而易舉但對演算法來說並不容易。

當模型處理 “it” 這個詞時, self-attention 會讓其和 “animal” 產生聯絡。模型逐個處理句中的單詞，self attention 將被處理詞和句中其他詞語關聯起來已獲得更好的編碼效果。

如果你對 RNN 比較熟悉，可以類比其中的 hidden state，它也讓 RNN 當前處理的內容可以用上先前處理過的詞語或向量的表示結果。Transformer 則是通過 Self-attention 機制來講關聯詞的語義資訊納入當前處理的詞中。

當編號為 #5 的編碼器（棧頂編碼器）開始編碼 “it” 一詞時，注意力機制的一部分落在了 “The Animal” 上，並把其表示結果納入 “it” 的編碼中。

你可以參考 Tensor2Tensor notebook 中的內容，以互動式視覺化的方式理解這部分內容。

Self-Attention 計算過程詳解

讓我們先學習如何向量來計算 self-attetion 結果，之後再看看如何利用矩陣應用它。

計算的第一步是為編碼器中的每個輸入向量（此處為每個詞語的嵌入結果）建立三個向量，分別為 Query 向量，Key 向量和 Value 向量。這些向量是通過將輸入向量和三個矩陣相乘得到的，而這三個矩陣也會在模型訓練過程中不斷更新。

注意這些向量的維度都比輸入的詞向量的維度要低，它們的長度為64，而輸入輸出向量維度為 512。這些向量也不是非得降維，這種選擇是為了讓計算更穩定。

將向量 x1 同 WQ 矩陣相乘會生成和輸入詞語關聯的 query 向量，同理 key 和 value 向量也是這麼生成的。它們是 attention 機制中所需的三類抽象，閱讀後續內容後你就會理解它們的作用機制。

self-attention 的第二步是計算 score 值，假定women在計算例子中的 “Thinking” 一詞，我們需要計算句中每個詞相對該詞的 score。score 的值代表了在編碼 “Thinking” 時該給其他詞分配多少注意力。

score 是通過將 query 向量 和 key 向量點乘得到的，當我們編碼句中第一個單詞時，第一個 score 由 q1 和 k1 點乘得到，第二個由 q1 和 k2 點乘得到。

第三步是將 score 除以 8（8 是 key 向量維數的平方根，這會使得梯度更加穩定，除以別的數字也可以，平方根是預設值），之後將結果進行 softmax，讓向量的元素和為1。

softmax 之後的 score 決定了每個詞對當前位置編碼的貢獻，該位置詞語本身顯然有最高的 score，但納入和當前詞相關的詞的資訊也是大有裨益的。

第五步是把每個 value 向量同 softmax socre 相乘作為賦權過程（為後續求和做準備），把值得注意的詞賦予高權重，無關緊要的詞賦予一個極小的權重。

第六步是將賦權後的 value 向量累加，這就是 self-attention 層在該位置的編碼輸出。

以上就是 self-attention 的計算全過程，輸出的結果矩陣將被輸入到前饋神經網路中。實際情況下整個計算過程是以矩陣計算的方式來提速的，下一節將闡述這個過程。

Self-Attention 的矩陣計算

現在的第一步是計算 Query， Key 和 Value 矩陣了，先把詞嵌入向量拼成矩陣，然後分別和 WQ， WK， WV 三個權重矩陣相乘。

、

矩陣 X 中的每一行為輸入的句子中的一個詞向量，我們可以清晰地看到輸入維度（512 維，圖中用 4 個方塊表示）和 q/k/v 向量（64 維，圖中為 3 個方塊）的差異。

最後，由於是矩陣計算，我們把 6 個計算步驟壓縮成一個來獲得 self-attention 層的結果。

多頭怪獸

論文後續有又將 self-attention 層融入到名為多頭注意力的機制中，從兩個層面提升了注意力機制的效果：

1. 它擴充了模型的注意力分配能力。上面的例子裡，z1 包含了所有其他詞編碼的資訊，但依舊是被自己位置的詞語的結果主導。當我們翻譯 “The animal didn’t cross the street because it was too tired” 時，我們希望知道 “it” 指代哪個詞。（讓 “it” 的編碼以其指代物件詞的資訊為主，譯者注）
2. 多頭扔注意力層有了更多的表示子空間。下文我們會看到在多頭注意力機制下，會有多組 Query/Key/Value 權重矩陣（Transformer 使用了 8 個注意力頭，所以對應也有 8 個編碼器和解碼器）。每組矩陣都是隨機初始化的，在模型訓練後，每組都用來將輸入的向量投影至表示子空間。

多頭注意力機制下我們將各個 WQ/WK/WV 矩陣分別儲存，依次讓輸入矩陣 X 與其計算生成每個頭的 Q/K/V 矩陣。

如果繼續採用之前的方式，就是重複 8 次計算流程，得到 8 個矩陣 Z。

這給我們帶來了一些些挑戰，後續的前饋層並不接受 8 個矩陣作為輸入，所以我們要把 8 個壓縮成一個。具體做法是把 Z 矩陣們橫向拼接，然後和一個權重矩陣 WO 相乘。

多頭注意力機制就這些內容，裡面涉及到了大量的矩陣，下圖我試著用一個圖來闡明他們的關係和作用。

瞭解多頭的機理後，我們回顧之前的例子，看看對 “it” 的編碼有什麼不同：

編碼過程中，一個注意力頭權重主要落在 “the animal” 上，另一個落在 “tired” 上，某種意義上， “it” 的表示結果包含了這兩個詞的資訊。

如果我們把所有注意力頭都放在一個圖裡，事情就變得更難解釋了：

用位置編碼表示劇中詞語的順序

目前我們一直沒有提到句中詞語的順序問題，為了表示它，transformer 在每個輸入的嵌入層增加了一個向量。這些向量用特定的模式表示詞語的順序或是詞語之間的距離。使得計算 Q/K/V 向量或是點乘生成注意力層時可以涵蓋距離資訊。

假定輸入的詞嵌入層維度為 4，實際的位置形式編碼如下（來自實際案例）：

那麼高維下的位置編碼什麼樣呢？

下圖中，每一行代表一個詞向量的位置編碼，第一行就是句子裡第一個詞的位置資訊，向量維度為 512，每個元素的取值咋吃 -1 到 1 之間，圖中用顏色表示數值大小。

一個由 20 個單詞構成的句子的真實案例，位置編碼維度為 512，可以看到圖從中間被一分為二，這是因為左右的值由不同的函式生成（ 左側為 sine，右側圍 cosine ）。雙邊的結果直接被拼接到了一起共同編碼位置資訊。

論文的 3.5 部分也詳細闡述了位置編碼的計算公式，程式碼中的 get_timing_signal_1d() 函式實現了這一過程。當然，位置編碼的生成方法遠不止這一種，但它的優勢在於可以編碼任意長度的句子。

殘差部分

在繼續講解之前還有一個細節需要闡明，即每個編碼器中的每一個子層（self-attention, 前饋神經網路）都有一個殘差連線結構，之後還有 layer-normalization](https://arxiv.org/abs/1607.06450) 步驟。

將向量和 self-attention 的 layer-norm 操作一起視覺化如下：

解碼器部分的 sub-layers 也有這樣的結構，下圖展示了一個包含兩個編碼器和解碼器的 Transformer 的結構：

解碼器

關於編碼器，我們已經說得夠多了，解碼器部分的元件和它也基本一致。但我們還需要了解一下它們是如何協同工作的。

編碼器的從輸入序列開始處理，最後輸出的是 attention 轉化得到的一系列向量 K 和 V。它們會在每個解碼器的 “encoder-decoder attention” 層被使用以輔助解碼器關注輸入序列的正確位置：

解碼步驟的每一步都生成輸出序列的一部分（本例中為英語句子）。

接下來的工作就是重複這一過程直到生成了告知解碼器工作結束的特定符號。而在下一個時間步，這些輸出都會餵給最底層的解碼器。就像編碼器一樣，解碼器的輸出也向上層層傳遞。位置編碼也一同輸入來表名每個詞語的位置。

解碼器中的 self-attention 層和編碼器中的稍稍有些不同，只納入當前輸出序列前的位置編碼資訊。這個限制通過在 softmax 處理前把後續位置編碼設為 -inf 實現。

“Encoder-Decoder Attention” 層和多頭 self-attention 層的工作方式大體一致，但它只從位於下方的層來生成 Query 矩陣，Key 和 Value 矩陣則用編碼器的輸出計算得到。

最後的 Linear 和 Softmax 層

編碼器棧會產出一個元素為浮點數的向量，怎麼把它轉換為一個詞呢？這就是 後續的 Linear 層和 Softmax 層的工作。

Linear 層是一個簡單的全連線網路，它將解碼器棧產出的向量投影成一個名為 logits 的大向量。

假設我們的模型從訓練集中學習了 10000 個不同的英文單詞（輸出詞典），那麼 logits 向量的長度就是 10000，每個元素的值是一個詞的概率得分。

而緊隨其後的 Softmax 層則把這些得分歸一化為概率值（全部為正數且和為1），概率最大的那個詞會被用作當前時間步的輸出。

上圖演示瞭解碼其棧輸出轉化為詞語的過程。

訓練過程回顧

我們已經瞭解了訓練好的 Transformer 是如何工作的了，現在就來看看它的訓練過程。

訓練過程中模型的處理過程和上文提到的一樣，但訓練集是有標註的，我們可以比較輸出和正確結果。為了方便圖示，我們假設輸出詞典只有 6 個詞（“a”, “am”, “i”, “thanks”, “student”, 和代表句子結尾的 “” ），對應索引如下：

之後英語句子就可以 one-hot 編碼了， 例如單詞 “am” 就可以表示為如下向量:

簡要回顧後, 就該看看損失函式了。

損失函式

假設我們現在正在訓練模型，樣本非常簡單 - 把 “merci” 翻譯為 “thanks”。我們希望最終輸出 “thanks” 的概率最大，但在模型沒有訓練好的情況下，這個情況一半不會發生。

由於模型的初始引數（權重）是隨機生成的，未訓練的模型輸出也是很隨機的。我們可以將其與實際結果比較，之後利用反向傳播來調節引數使得輸出和標註靠攏。那麼問題來了，怎麼比較兩個離散概率分佈呢，我們可以簡單地採用相減的方式，更多細節可以參考 cross-entropy 和 Kullback–Leibler divergence。

我們現在使用的是高度簡化的例子，實際情況中一個句子往往不止一個單詞。同時我們希望訓練好的模型滿足以下性質：

• 使用和詞典詞語數等長的向量表示每一個位置的概率分佈（例子裡為6，實際上會是 30000 或 50000 這麼大的數字）
• 第一個位置的概率分佈中單詞 “i” 對應的概率值最大
• 第一個位置的概率分佈中單詞 “am” 對應的概率值最大
• And so on, until the fifth output distribution indicates ‘<end of sentence>’ symbol, which also has a cell associated with it from the 10,000 element vocabulary.以此類推，直到第五個位置輸出 “<end of sentence>”

上圖為訓練集標註轉換為向量的結果。

當我們在一個足夠大的資料集上訓練足夠長的時間之後，我們期望輸出的概率分佈應當接近下圖所示：

我們希望模型在經過訓練之後能輸出我們期望的結果，當然並不是指只在訓練集上能很好的擬合( 詳見: cross validation )。請注意，不論是哪個位置，每個單詞相應的概率都不為 0，哪怕這本身是個徹底的不可能事件。這是 softmax 函式的性質決定的，它對於模型訓練很有幫助。

目前模型的是把每個位置概率最高的詞語輸出，這種方式稱為貪心解碼（ greedy decoding )。另一種方式則是保留每個位置概率最大的兩個詞（比如 “I” 和 “a”），下一個位置則假定前一個詞語分別為二者跑兩次模型，同樣保留概率最大的兩個結果，在對後續位置採用相同策略。這個方法叫做集束搜尋 ( beam search )。在我們的例子裡搜尋的範圍為 2 （即每個時間步會保留概率最大的前兩個路徑），這是我們可以通過實驗效果來確定超引數。

繼續深入

我希望本文能讓你建立起對 transformer 的初步瞭解，如果你想更深入的研究，我建議如下方式：

• 閱讀 Attention Is All You Need 原文, 相關部落格 (Transformer: A Novel Neural Network Architecture for Language Understanding), 和 Tensor2Tensor announcement.
• 觀看 Łukasz Kaiser 的演講 瞭解模型細節
• 試用 Jupyter Notebook provided as part of the Tensor2Tensor repo
• 瀏覽 Tensor2Tensor repo.

之後的工作:

• Depthwise Separable Convolutions for Neural Machine Translation
• One Model To Learn Them All
• Discrete Autoencoders for Sequence Models
• Generating Wikipedia by Summarizing Long Sequences
• Image Transformer
• Training Tips for the Transformer Model
• Self-Attention with Relative Position Representations
• Fast Decoding in Sequence Models using Discrete Latent Variables
• Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

感謝

Thanks to Illia Polosukhin, Jakob Uszkoreit, Llion Jones , Lukasz Kaiser, Niki Parmar, and Noam Shazeer for providing feedback on earlier versions of this post.