深入理解Transformer及其源碼解讀

2019 年 10 月 24 日
筆記

　　深度學習廣泛應用於各個領域。基於transformer的預訓練模型（gpt/bertd等）基本已統治NLP深度學習領域，可見transformer的重要性。本文結合《Attention is all you need》Harvard 的程式碼《Annotated Transformer》深入理解transformer模型。 Harvard的程式碼在python3.6 torch 1.0.1 上跑不通，本文做了很多修改。修改後的程式碼地址：Transformer。

1 模型的思想

　　Transformer中拋棄了傳統的CNN和RNN，整個網路結構完全是由Attention機制組成。作者採用Attention機制的原因是考慮到RNN（或者LSTM，GRU等）的計算限制為是順序的，也就是說RNN相關演算法只能從左向右依次計算或者從右向左依次計算，這種機制帶來了兩個問題：　

　　(1) 時間片 $t$ 的計算依賴 $t-1$ 時刻的計算結果，這樣限制了模型的並行能力；

　　(2) 順序計算的過程中資訊會丟失，儘管LSTM等門機制的結構一定程度上緩解了長期依賴的問題，但是對於特別長期的依賴現象，LSTM依舊無能為力。

　　Transformer的提出解決了上面兩個問題：

　　(1) 首先它使用了Attention機制，將序列中的任意兩個位置之間的距離是縮小為一個常量；

　　(2) 其次它不是類似RNN的順序結構，因此具有更好的並行性，符合現有的GPU框架。

2 模型的架構

　　如上圖，transformer模型本質上是一個Encoder-Decoder的結構。輸入序列先進行Embedding，經過Encoder之後結合上一次output再輸入Decoder，最後用softmax計算序列下一個單詞的概率。

3 Embedding

　　transformer的輸入是Word Embedding + Position Embedding。

3.1 Word Embedding

　　Word embedding在pytorch中通常用 nn.Embedding 實現，其權重矩陣通常有兩種選擇：

　　（1）使用 Pre-trained的Embeddings並固化，這種情況下實際就是一個 Lookup Table。

　　（2）對其進行隨機初始化(當然也可以選擇 Pre-trained 的結果)，但設為 Trainable。這樣在 training 過程中不斷地對 Embeddings 進行改進。

　　transformer選擇後者，程式碼實現如下：

class Embeddings(nn.Module):      def __init__(self, d_model, vocab):          super(Embeddings, self).__init__()          self.lut = nn.Embedding(vocab, d_model)          self.d_model = d_model  #表示embedding的維度        def forward(self, x):          return self.lut(x) * math.sqrt(self.d_model)

　　其中d_model表示embedding的維度，即詞向量的維度；vocab表示辭彙表的數量。

3.2 Positional Embedding

　　在RNN中，對句子的處理是一個個word按順序輸入的。但在 Transformer 中，輸入句子的所有word是同時處理的，沒有考慮詞的排序和位置資訊。因此，Transformer 的作者提出了加入 “positional encoding” 的方法來解決這個問題。“positional encoding“”使得 Transformer 可以衡量 word 位置有關的資訊。

　　如何實現具有位置資訊的encoding呢？作者提供了兩種思路：

通過訓練學習 positional encoding 向量；
使用公式來計算 positional encoding向量。

　　試驗後發現兩種選擇的結果是相似的，所以採用了第2種方法，優點是不需要訓練參數，而且即使在訓練集中沒有出現過的句子長度上也能用。

　　Positional Encoding的公式如下：

$$PE_{(pos,2i)} = sin(pos / 10000^{2i/d_{text{model}}})$$

$$PE_{(pos,2i+1)} = cos(pos / 10000^{2i/d_{text{model}}})$$

　　其中，$pos$指的是這個 word 在這個句子中的位置；$2i$指的是 embedding 詞向量的偶數維度，$2i+1$指的是embedding 詞向量的奇數維度。

具體實現如下：

# Positional Encoding  class PositionalEncoding(nn.Module):      "實現PE功能"      def __init__(self, d_model, dropout, max_len=5000):          super(PositionalEncoding, self).__init__()          self.dropout = nn.Dropout(p=dropout)            pe = torch.zeros(max_len, d_model)          position = torch.arange(0., max_len).unsqueeze(1)          div_term = torch.exp(torch.arange(0., d_model, 2) *                               -(math.log(10000.0) / d_model))            pe[:, 0::2] = torch.sin(position * div_term)    # 偶數列          pe[:, 1::2] = torch.cos(position * div_term)    # 奇數列          pe = pe.unsqueeze(0)           # [1, max_len, d_model]          self.register_buffer('pe', pe)        def forward(self, x):          x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)          return self.dropout(x)

　　注意：“x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)” 這行程式碼表示；輸入模型的整個Embedding是Word Embedding與Positional Embedding直接相加之後的結果。

　　為什麼上面的兩個公式能體現單詞的相對位置資訊呢？

　　我們寫一段程式碼取詞向量的4個維度看下：

# 在位置編碼下方，將基於位置添加正弦波。對於每個維度，波的頻率和偏移都不同。  plt.figure(figsize=(15, 5))  pe = PositionalEncoding(20, 0)  y = pe.forward(Variable(torch.zeros(1, 100, 20)))  plt.plot(np.arange(100), y[0, :, 4:8].data.numpy())  plt.legend(["dim %d"%p for p in [4,5,6,7]])

　　輸出影像：

　　可以看到某個序列中不同位置的單詞，在某一維度上的位置編碼數值不一樣，即同一序列的不同單詞在單個緯度符合某個正弦或者餘弦，可認為他們的具有相對關係。

4 Encoder

　　Encoder部分是由個層相同小Encoder Layer串聯而成。小Encoder Layer可以簡化為兩個部分：（1）Multi-Head Self Attention (2) Feed-Forward network。示意圖如下:

　　事實上multi head self attention 和feed forward network之後都接了一層add 和norm這裡先不講，後面4.1.2再講。

4.1 Muti-Head-Attention

　　Multi-Head Self Attention 實際上是由h個Self Attention 層並行組成，原文中h=8。接下來我們先介紹Self Attention。

4.1.1 Self-Attention

　　self-attention的輸入是序列詞向量，此處記為x。x經過一個線性變換得到query(Q), x經過第二個線性變換得到key(K), x經過第三個線性變換得到value(V)。

也就是：

key = linear_k(x)
query = linear_q(x)
value = linear_v(x)

用矩陣表示即：

　　注意：這裡的linear_k, linear_q, linear_v是相互獨立、權重（$W^Q$, $W^K$, $W^V$)是不同的，通過訓練可得到。得到query(Q)，key(K)，value(V)之後按照下面的公式計算attention(Q, K, V)：

$$Attention(Q, K, V) = Softmax(frac{QK^T}{sqrt{d_k}})V$$

用矩陣表示上面的公式即：

　　這裡Z就是attention(Q, K, V)。

　　(1) 這裡$d_k=d_{model}/h = 512/8 = 64$。

　　(2) 為什麼要用$sqrt{d_k}$ 對 $QK^T$進行縮放呢？

　　$d_k$實際上是Q/K/V的最後一個維度，當$d_k$越大，$QK^T$就越大，可能會將softmax函數推入梯度極小的區域。

　　(3) softmax之後值都介於0到1之間，可以理解成得到了 attention weights。然後基於這個 attention weights 對 V 求 weighted sum 值 Attention(Q, K, V)。

　　Multi-Head-Attention 就是將embedding之後的X按維度$d_{model}=512$ 切割成$h=8$個，分別做self-attention之後再合併在一起。

源碼如下：

class MultiHeadedAttention(nn.Module):      def __init__(self, h, d_model, dropout=0.1):          "Take in model size and number of heads."          super(MultiHeadedAttention, self).__init__()          assert d_model % h == 0          self.d_k = d_model // h          self.h = h          self.linears = clones(nn.Linear(d_model, d_model), 4)          self.attn = None          self.dropout = nn.Dropout(p=dropout)        def forward(self, query, key, value, mask=None):          """          實現MultiHeadedAttention。             輸入的q，k，v是形狀 [batch, L, d_model]。             輸出的x 的形狀同上。          """          if mask is not None:              # Same mask applied to all h heads.              mask = mask.unsqueeze(1)          nbatches = query.size(0)            # 1) 這一步qkv變化:[batch, L, d_model] ->[batch, h, L, d_model/h]          query, key, value =               [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)                     for l, x in zip(self.linears, (query, key, value))]            # 2) 計算注意力attn 得到attn*v 與attn          # qkv :[batch, h, L, d_model/h] -->x:[b, h, L, d_model/h], attn[b, h, L, L]          x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)          # 3) 上一步的結果合併在一起還原成原始輸入序列的形狀          x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)          # 最後再過一個線性層          return self.linears[-1](x)

4.1.2 Add & Norm

　　x 序列經過multi-head-self-attention 之後實際經過一個“add+norm”層，再進入feed-forward network(後面簡稱FFN)，在FFN之後又經過一個norm再輸入下一個encoder layer。

class LayerNorm(nn.Module):      """構造一個layernorm模組"""      def __init__(self, features, eps=1e-6):          super(LayerNorm, self).__init__()          self.a_2 = nn.Parameter(torch.ones(features))          self.b_2 = nn.Parameter(torch.zeros(features))          self.eps = eps        def forward(self, x):          "Norm"          mean = x.mean(-1, keepdim=True)          std = x.std(-1, keepdim=True)          return self.a_2 * (x - mean) / (std + self.eps) + self.b_2      class SublayerConnection(nn.Module):      """Add+Norm"""      def __init__(self, size, dropout):          super(SublayerConnection, self).__init__()          self.norm = LayerNorm(size)          self.dropout = nn.Dropout(dropout)        def forward(self, x, sublayer):          "add norm"          return x + self.dropout(sublayer(self.norm(x)))

　　注意：幾乎每個sub layer之後都會經過一個歸一化，然後再加在原來的輸入上。這裡叫殘餘連接。

4.2 Feed-Forward Network

　　Feed-Forward Network可以細分為有兩層，第一層是一個線性激活函數，第二層是激活函數是ReLU。可以表示為：

$$FFN=max(0, xW_1+b_1)W_2 + b_2$$

　　這層比較簡單，就是實現上面的公式，直接看程式碼吧：

# Position-wise Feed-Forward Networks  class PositionwiseFeedForward(nn.Module):      "實現FFN函數"      def __init__(self, d_model, d_ff, dropout=0.1):          super(PositionwiseFeedForward, self).__init__()          self.w_1 = nn.Linear(d_model, d_ff)          self.w_2 = nn.Linear(d_ff, d_model)          self.dropout = nn.Dropout(dropout)        def forward(self, x):          return self.w_2(self.dropout(F.relu(self.w_1(x))))

　　總的來說Encoder 是由上述小encoder layer 6個串列疊加組成。encoder sub layer主要包含兩個部分：

SubLayer-1 做 Multi-Headed Attention
SubLayer-2 做 Feed Forward Neural Network

　　來看下Encoder主架構的程式碼：

def clones(module, N):      "產生N個相同的層"      return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])    class Encoder(nn.Module):      """N層堆疊的Encoder"""      def __init__(self, layer, N):          super(Encoder, self).__init__()          self.layers = clones(layer, N)          self.norm = LayerNorm(layer.size)        def forward(self, x, mask):          "每層layer依次通過輸入序列與mask"          for layer in self.layers:              x = layer(x, mask)          return self.norm(x)

5 Decoder

　　Decoder與Encoder有所不同，Encoder與Decoder的關係可以用下圖描述（以機器翻譯為例）：

Decoder的程式碼主要結構：

# Decoder部分  class Decoder(nn.Module):      """帶mask功能的通用Decoder結構"""      def __init__(self, layer, N):          super(Decoder, self).__init__()          self.layers = clones(layer, N)          self.norm = LayerNorm(layer.size)        def forward(self, x, memory, src_mask, tgt_mask):          for layer in self.layers:              x = layer(x, memory, src_mask, tgt_mask)          return self.norm(x)

Decoder子結構（Sub layer）：

　　Decoder 也是N=6層堆疊的結構。被分為3個 SubLayer，Encoder與Decoder有三大主要的不同：

　　（1）Decoder SubLayer-1 使用的是 “Masked” Multi-Headed Attention 機制，防止為了模型看到要預測的數據，防止泄露。

　　（2）SubLayer-2 是一個 Encoder-Decoder Multi-head Attention。

　　(3) LinearLayer 和 SoftmaxLayer 作用於 SubLayer-3 的輸出後面，來預測對應的 word 的 probabilities 。

5.1 Mask-Multi-Head-Attention

　　Mask 的目的是防止 Decoder “seeing the future”，就像防止考生偷看考試答案一樣。這裡mask是一個下三角矩陣，對角線以及對角線左下都是1，其餘都是0。下面是個10維度的下三角矩陣：

tensor([[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],           [1, 1, 0, 0, 0, 0, 0, 0, 0, 0],           [1, 1, 1, 0, 0, 0, 0, 0, 0, 0],           [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],           [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],           [1, 1, 1, 1, 1, 1, 0, 0, 0, 0],           [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],           [1, 1, 1, 1, 1, 1, 1, 1, 0, 0],           [1, 1, 1, 1, 1, 1, 1, 1, 1, 0],           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]], dtype=torch.uint8)

Mask的程式碼實現：

def subsequent_mask(size):      """      mask後續的位置，返回[size, size]尺寸下三角Tensor      對角線及其左下角全是1，右上角全是0      """      attn_shape = (1, size, size)      subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')      return torch.from_numpy(subsequent_mask) == 0

　　當mask不為空的時候，attention計算需要將x做一個操作：scores = scores.masked_fill(mask == 0, -1e9)。即將mask==0的替換為-1e9,其餘不變。

5.2 Encoder-Decoder Multi-head Attention

　　這部分和Multi-head Attention的區別是該層的輸入來自encoder和上一次decoder的結果。具體實現如下：

class DecoderLayer(nn.Module):      "Decoder is made of self-attn, src-attn, and feed forward (defined below)"      def __init__(self, size, self_attn, src_attn, feed_forward, dropout):          super(DecoderLayer, self).__init__()          self.size = size          self.self_attn = self_attn          self.src_attn = src_attn          self.feed_forward = feed_forward          self.sublayer = clones(SublayerConnection(size, dropout), 3)        def forward(self, x, memory, src_mask, tgt_mask):          "將decoder的三個Sublayer串聯起來"          m = memory          x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))          x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))          return self.sublayer[2](x, self.feed_forward)

　　注意：self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask)) 這行就是Encoder-Decoder Multi-head Attention。

　　query = x，key = m, value = m, mask = src_mask，這裡x來自上一個 DecoderLayer，m來自 Encoder的輸出。

5.3 Linear and Softmax to Produce Output Probabilities

　　Decoder的最後一個部分是過一個linear layer將decoder的輸出擴展到與vocabulary size一樣的維度上。經過softmax 後，選擇概率最高的一個word作為預測結果。假設我們有一個已經訓練好的網路，在做預測時，步驟如下：

　　（1）給 decoder 輸入 encoder 對整個句子 embedding 的結果和一個特殊的開始符號 </s>。decoder 將產生預測，在我們的例子中應該是 ”I”。

　　（2）給 decoder 輸入 encoder 的 embedding 結果和 “</s>I”，在這一步 decoder 應該產生預測 “am”。

　　（3）給 decoder 輸入 encoder 的 embedding 結果和 “</s>I am”，在這一步 decoder 應該產生預測 “a”。

　　（4）給 decoder 輸入 encoder 的 embedding 結果和 “</s>I am a”，在這一步 decoder 應該產生預測 “student”。

　　（5）給 decoder 輸入 encoder 的 embedding 結果和 “</s>I am a student”, decoder應該生成句子結尾的標記，decoder 應該輸出 ”</eos>”。

　　（6）然後 decoder 生成了 </eos>，翻譯完成。
　　這部分的程式碼實現：

class Generator(nn.Module):      """      Define standard linear + softmax generation step。      定義標準的linear + softmax 生成步驟。      """      def __init__(self, d_model, vocab):          super(Generator, self).__init__()          self.proj = nn.Linear(d_model, vocab)        def forward(self, x):          return F.log_softmax(self.proj(x), dim=-1)

　　在訓練過程中，模型沒有收斂得很好時，Decoder預測產生的詞很可能不是我們想要的。這個時候如果再把錯誤的數據再輸給Decoder，就會越跑越偏。這個時候怎麼辦？

　（1）在訓練過程中可以使用 “teacher forcing”。因為我們知道應該預測的word是什麼，那麼可以給Decoder喂一個正確的結果作為輸入。

　　(2）除了選擇最高概率的詞 (greedy search)，還可以選擇是比如 “Beam Search”，可以保留topK個預測的word。 Beam Search 方法不再是只得到一個輸出放到下一步去訓練了，我們可以設定一個值，拿多個值放到下一步去訓練，這條路徑的概率等於每一步輸出的概率的乘積。

6 Transformer的優缺點

6.1 優點

　　（1）每層計算複雜度比RNN要低。

　　（2）可以進行並行計算。

　　（3）從計算一個序列長度為n的資訊要經過的路徑長度來看, CNN需要增加卷積層數來擴大視野，RNN需要從1到n逐個進行計算，而Self-attention只需要一步矩陣計算就可以。Self-Attention可以比RNN更好地解決長時依賴問題。當然如果計算量太大，比如序列長度N大於序列維度D這種情況，也可以用窗口限制Self-Attention的計算數量。

　　（4）從作者在附錄中給出的栗子可以看出，Self-Attention模型更可解釋，Attention結果的分布表明了該模型學習到了一些語法和語義資訊。

6.2 缺點

　　在原文中沒有提到缺點，是後來在Universal Transformers中指出的，主要是兩點：

　　（1）實踐上：有些RNN輕易可以解決的問題transformer沒做到，比如複製string，或者推理時碰到的sequence長度比訓練時更長（因為碰到了沒見過的position embedding）。

　　（2）理論上：transformers不是computationally universal(圖靈完備)，這種非RNN式的模型是非圖靈完備的的，無法單獨完成NLP中推理、決策等計算問題（包括使用transformer的bert模型等等）。

7 References

　　1 http://jalammar.github.io/illustrated-transformer/

　　2 https://zhuanlan.zhihu.com/p/48508221

　　3 https://zhuanlan.zhihu.com/p/47063917

　　4 https://zhuanlan.zhihu.com/p/80986272

　　5 https://arxiv.org/abs/1706.03762