基於注意力機制，機器之心帶你理解與訓練神經機器翻譯系統

來自專欄機器之心64 人贊了文章

本文是機器之心 GitHub 實現項目，我們根據谷歌的 Transformer 原論文與 Harvard NLP 所實現的代碼學習構建了一個神經機器翻譯系統。因此，我們希望各位讀者也能根據這篇文章了解 Transformer 的架構，並動手實現一個神經機器翻譯系統。

自去年 6 月份「Attention is All You Need」發表以來，Transformer 受到越來越多的關注。它除了能顯著提升翻譯質量，同時還為很多 NLP 任務提供了新的架構。這篇論文放棄了傳統基於 RNN 或 CNN 的深度架構，並只保留了注意力（Attentaion）機制，雖然原論文在這一方面描述地比較清楚，但要正確地實現這樣的新型架構可能非常困難。

在這篇文章中，我們從注意力機制到神經機器翻譯系統解釋了實現 Transformer 的架構與代碼，並藉助這些實現理解原論文。機器之心整理了整個實現，並根據我們對原論文與實現的理解添加一些解釋。整個文章就是一個可運行的 Jupyter Notebook，讀者可直接在 Colaboratory 中閱讀文章與運行代碼。

• 機器之心實現地址：https://github.com/jiqizhixin/ML-Tutorial-Experiment
• 原實現地址：https://github.com/harvardnlp/annotated-transformer

本文所有的代碼都可以在谷歌 Colab 上運行，且讀者也可以在 GitHub 中下載全部的代碼在本地運行。這篇文章非常適合於研究者與感興趣的開發者，代碼很大程度上都依賴於 OpenNMT 庫。

在運行模型前，我們需要確保有對應的環境。如果在本地運行，那麼需要確保以下基本庫的導入不會報錯，若在 Colab 上運行，那麼首先需要運行以下第一個 pip 語句安裝對應的包。Colab 的環境配置非常簡單，一般只需要使用 conda 或 pip 命令就能完成。此外，Colab 語句前面加上「!」表示這是命令行，而不加感嘆號則表示這個代碼框是 Python 代碼。

# !pip install http://download.pytorch.org/whl/cu80/torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl numpy matplotlib spacy torchtext seaborn import numpy as npimport torchimport torch.nn as nnimport torch.nn.functional as Fimport math, copy, timefrom torch.autograd import Variableimport matplotlib.pyplot as pltimport seabornseaborn.set_context(context="talk")%matplotlib inline

引言

減少序列計算的任務目標構成了 Extended Neural GPU、ByteNet 和 ConvS2S 的基礎，它們都是使用卷積神經網路作為基本構建塊，因而能對所有輸入與輸出位置的隱藏表徵執行並行計算。在這些模型中，兩個任意輸入與輸出位置的信號關聯所需要的運算數量與它們的位置距離成正比，對於 ConvS2S 為線性增長，對於 ByteNet 為對數增長。這種現象使得學習較遠位置的依賴關係非常困難。而在 Transformer 中，這種成本會減少到一個固定的運算數量，儘管平均注意力位置加權會減少有效表徵力，但使用 Multi-Head Attention 注意力機制可以抵消這種成本。

自注意力（Self-attention），有時也稱為內部注意力，它是一種涉及單序列不同位置的注意力機制，並能計算序列的表徵。自注意力在多種任務中都有非常成功的應用，例如閱讀理解、摘要概括、文字蘊含和語句表徵等。自注意力這種在序列內部執行 Attention 的方法可以視為搜索序列內部的隱藏關係，這種內部關係對於翻譯以及序列任務的性能非常重要。

然而就我們所知道的，Transformer 是第一種完全依賴於自注意力以計算輸入與輸出表徵的方法，這意味著它沒有使用序列對齊的 RNN 或卷積網路。從 Transformer 的結構就可以看出，它並沒有使用深度網路抽取序列特徵，頂多使用幾個線性變換對特徵進行變換。

本文主要從模型架構、訓練配置和兩個實際翻譯模型開始介紹 Ashish Vaswani 等人的原論文與 Harvard NLP 團隊實現的代碼。在模型架構中，我們將討論編碼器、解碼器、注意力機制以及位置編碼等關鍵組成部分，而訓練配置將討論如何抽取批量數據、設定訓練循環、選擇最優化方法和正則化器等。最後我們將跟隨 Alexander Rush 等人的實現訓練兩個神經機器翻譯系統，其中一個僅使用簡單的合成數據，而另一個則是真實的 IWSLT 德語-英語翻譯數據集。

模型架構

大多數神經序列模型都使用編碼器-解碼器框架，其中編碼器將表徵符號的輸入序列 (x_1, …, x_n) 映射到連續表徵 z=(z_1, …, z_n)。給定中間變數 z，解碼器將會生成一個輸出序列 (y_1,…,y_m)。在每一個時間步上，模型都是自回歸的（auto-regressive），當生成序列中的下一個元素時，先前生成的元素會作為輸入。

以下展示了一個標準的編碼器-解碼器框架，EncoderDecoder 類定義了先編碼後解碼的過程，例如先將英文序列編碼為一個隱向量，在基於這個中間表徵解碼為中文序列。

class EncoderDecoder(nn.Module): """ A standard Encoder-Decoder architecture. Base for this and many other models. """ def __init__(self, encoder, decoder, src_embed, tgt_embed, generator): super(EncoderDecoder, self).__init__() self.encoder = encoder self.decoder = decoder self.src_embed = src_embed self.tgt_embed = tgt_embed self.generator = generator def forward(self, src, tgt, src_mask, tgt_mask): "Take in and process masked src and target sequences." return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask) def encode(self, src, src_mask): return self.encoder(self.src_embed(src), src_mask) def decode(self, memory, src_mask, tgt, tgt_mask): return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)class Generator(nn.Module): "Define standard linear + softmax generation step." 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)

Transformer 的整體架構也採用了這種編碼器-解碼器的框架，它使用了多層自注意力機制和層級歸一化，編碼器和解碼器都會使用全連接層和殘差連接。Transformer 的整體結構如下圖所示：

如上所示，左側為輸入序列的編碼器。輸入序列首先會轉換為詞嵌入向量，在與位置編碼向量相加後可作為 Multi-Head Attention 模塊的輸入，該模塊的輸出在與輸入相加後將投入層級歸一化函數，得出的輸出在饋送到全連接層後可得出編碼器模塊的輸出。這樣相同的 6 個編碼器模塊（N=6）可構成整個編碼器架構。解碼器模塊首先同樣構建了一個自注意力模塊，然後再結合編碼器的輸出實現 Multi-Head Attention，最後投入全連接網路並輸出預測詞概率。

這裡只是簡單地介紹了模型的大概過程，很多如位置編碼、Multi-Head Attention 模塊、層級歸一化、殘差鏈接和逐位置前饋網路等概念都需要讀者詳細閱讀下文，最後再回過頭理解完整的過程。

編碼器與解碼器堆棧

• 編碼器

編碼器由相同的 6 個模塊堆疊而成，每一個模塊都有兩個子層級構成。其中第一個子層級是 Multi-Head 自注意機制，其中自注意力表示輸入和輸出序列都是同一條。第二個子層級採用了全連接網路，主要作用在於注意子層級的特徵。此外，每一個子層級都會添加一個殘差連接和層級歸一化。

以下定義了編碼器的主體框架，在 Encoder 類中，每一個 layer 表示一個編碼器模塊，這個編碼器模塊由兩個子層級組成。layer 函數的輸出表示經過層級歸一化的編碼器模塊輸出，通過 For 循環堆疊層級就能完成整個編碼器的構建。

def clones(module, N): "Produce N identical layers." return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])class Encoder(nn.Module): "Core encoder is a stack of N layers" def __init__(self, layer, N): super(Encoder, self).__init__() self.layers = clones(layer, N) self.norm = LayerNorm(layer.size) def forward(self, x, mask): "Pass the input (and mask) through each layer in turn." for layer in self.layers: x = layer(x, mask) return self.norm(x)

如編碼器的結構圖所示，每個子層級都會會添加一個殘差連接，並隨後傳入層級歸一化。上面構建的主體架構也調用了層級歸一化函數，以下代碼展示了層級歸一化的定義。

class LayerNorm(nn.Module): "Construct a layernorm module (See citation for details)." 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): mean = x.mean(-1, keepdim=True) std = x.std(-1, keepdim=True) return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

層級歸一化可以通過修正每一層內激活值的均值與方差而大大減少協方差偏離問題。簡單來說，一個層級的均值可以通過計算該層所有神經元激活值的平均值而得出，然後再根據均值計算該層所有神經元激活值的方差。最後根據均值與方差，我們可以對這一層所有輸出值進行歸一化。

如上 LayerNorm 類所示，我們首先需要使用方法 mean 求輸入 x 最後一個維度的均值，keepdim 為真表示求均值後的維度保持不變，並且均值會廣播操作到對應的維度。同樣使用 std 方法計算標準差後，該層所有激活值分別減去均值再除以標準差就能實現歸一化，分母加上一個小值 eps 可以防止分母為零。

因此，每一個子層的輸出為 LayerNorm(x+Sublayer(x))，其中 Sublayer(x) 表示由子層本身實現的函數。我們應用 Dropout 將每一個子層的輸出隨機失活，這一過程會在加上子層輸入和執行歸一化之前完成。

以下定義了殘差連接，我們會在投入層級歸一化函數前將子層級的輸入與輸出相加。為了使用這些殘差連接，模型中所有的子層和嵌入層的輸出維度都是 d_model=512。

class SublayerConnection(nn.Module): """ A residual connection followed by a layer norm. Note for code simplicity the norm is first as opposed to last. """ def __init__(self, size, dropout): super(SublayerConnection, self).__init__() self.norm = LayerNorm(size) self.dropout = nn.Dropout(dropout) def forward(self, x, sublayer): "Apply residual connection to any sublayer with the same size." return x + self.dropout(sublayer(self.norm(x)))

在上述代碼定義中，x 表示上一層添加了殘差連接的輸出，這一層添加了殘差連接的輸出需要將 x 執行層級歸一化，然後饋送到 Multi-Head Attention 層或全連接層，添加 Dropout 操作後可作為這一子層級的輸出。最後將該子層的輸出向量與輸入向量相加得到下一層的輸入。

編碼器每個模塊有兩個子層，第一個為 multi-head 自注意力層，第二個為簡單的逐位置全連接前饋網路。以下的 EncoderLayer 類定義了一個編碼器模塊的過程。

class EncoderLayer(nn.Module): "Encoder is made up of self-attn and feed forward (defined below)" def __init__(self, size, self_attn, feed_forward, dropout): super(EncoderLayer, self).__init__() self.self_attn = self_attn self.feed_forward = feed_forward self.sublayer = clones(SublayerConnection(size, dropout), 2) self.size = size def forward(self, x, mask): "Follow Figure 1 (left) for connections." x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) return self.sublayer[1](x, self.feed_forward)

以上代碼疊加了自注意力層與全連接層，其中 Multi-Head Attention 機制的輸入 Query、Key 和 Value 都為 x 就表示自注意力。

• 解碼器

解碼器也由相同的 6 個模塊堆疊而成，每一個解碼器模塊都有三個子層組成，每一個子層同樣會加上殘差連接與層級歸一化運算。第一個和第三個子層分別與編碼器的 Multi-Head 自注意力層和全連接層相同，而第二個子層所採用的 Multi-Head Attention 機制使用編碼器的輸出作為 Key 和 Value，採用解碼模塊第一個子層的輸出作為 Query。

我們同樣需要修正編碼器堆棧中的自注意力子層，以防止當前位置注意到後續序列位置，這一修正可通過掩碼實現。以下的解碼器的主體堆疊結構和編碼器相似，只需要簡單地堆疊解碼器模塊就能完成。

class Decoder(nn.Module): "Generic N layer decoder with masking." 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)

以下展示了一個解碼器模塊的架構，第一個 Multi-Head Attention 機制的三個輸入都是 x，因此它是自注意力。第二個 Multi-Head 注意力機制輸入的 Key 和 Value 是編碼器的輸出 memory，輸入的 Query 是上一個子層的輸出 x。最後在疊加一個全連接網路以完成一個編碼器模塊的構建。

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): "Follow Figure 1 (right) for connections." 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)

對於序列建模來說，模型應該只能查看有限的序列信息。例如在時間步 i，模型能讀取整個輸入序列，但只能查看時間步 i 及之前的序列信息。對於 Transformer 的解碼器來說，它會輸入整個目標序列，且注意力機制會注意到整個目標序列各個位置的信息，因此我們需要限制注意力機制能看到的信息。

如上所述，Transformer 在注意力機制中使用 subsequent_mask 函數以避免當前位置注意到後面位置的信息。因為輸出詞嵌入是位置的一個偏移，因此我們可以確保位置 i 的預測僅取決於在位置 i 之前的已知輸出。

def subsequent_mask(size): "Mask out subsequent positions." attn_shape = (1, size, size) subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype(uint8) return torch.from_numpy(subsequent_mask) == 0

以下為注意力掩碼的可視化，其中每一行為一個詞，每一列則表示一個位置。下圖展示了每一個詞允許查看的位置，訓練中詞是不能注意到未來詞的。

plt.figure(figsize=(5,5))plt.imshow(subsequent_mask(20)[0])None

注意力機制

谷歌在原論文中展示了注意力機制的一般化定義，即它和 RNN 或 CNN 一樣也是一種編碼序列的方案。一個注意力函數可以描述為將 Query 與一組鍵值對（Key-Value）映射到輸出，其中 Query、Key、Value 和輸出都是向量。輸出可以通過值的加權和而計算得出，其中分配到每一個值的權重可通過 Query 和對應 Key 的適應度函數（compatibility function）計算。

在翻譯任務中，Query 可以視為原語詞向量序列，而 Key 和 Value 可以視為目標語詞向量序列。一般的注意力機制可解釋為計算 Query 和 Key 之間的相似性，並利用這種相似性確定 Query 和 Value 之間的注意力關係。

以下是點積注意力的結構示意圖，我們稱這種特殊的結構為「縮放點積注意力」。它的輸入由維度是 d_k 的 Query 和 Key 組成，Value 的維度是 d_v。如下所示，我們會先計算 Query 和所有 Key 的點乘，並每一個都除上 squre_root(d_k) 以防止乘積結果過大，然後再饋送到 Softmax 函數以獲得與 Value 對應的權重。根據這樣的權重，我們就可以配置 Value 向量而得出最後的輸出。

Image(filename=images/ModalNet-19.png

在上圖中，Q 和 K 的運算有一個可選的 Mask 過程。在編碼器中，我們不需要使用它限制注意力模塊所關注的序列信息。而在解碼器中，我們需要它限制注意力模塊只能注意到當前時間步及之前時間步的信息。這一個過程可以很簡潔地表示為函數 Attention(Q, K, V)。

Attention(Q, K, V) 函數在輸入矩陣 Q、K 和 V 的情況下可計算 Query 序列與 Value 序列之間的注意力關係。其中 Q 的維度為 n×d_k，表示有 n 條維度為 d_k 的 Query、K 的維度為 m×d_k、V 的維度為 m×d_v。這三個矩陣的乘積可得出 n×d_v 維的矩陣，它表示 n 條 Query 對應注意到的 Value 向量。

原論文作者發現，當每一條 Query 的維度 d_k 比較小時，點乘注意力和加性注意力的性能相似，但隨著 d_k 的增大，加性注意力的性能會超過點乘注意力機制。不過點乘注意力有一個強大的屬性，即它可以利用矩陣乘法的並行運算大大加快訓練速度。

原論文作者認為點乘注意力效果不好的原因是在 d_k 比較大的情況下，乘積結果會非常大，因此會導致 Softmax 快速飽和並只能提供非常小的梯度來更新參數。所以他們採用了根號下 d_k 來縮小點乘結果，並防止 Softmax 函數飽和。

為了證明為什麼點積的量級會變得很大，我們假設元素 q 和 k 都是均值為 0、方差為 1 的獨立隨機變數，它們的點乘 q?k=∑q_i*k_i 有 0 均值和 d_k 的方差。為了抵消這種影響，我們可以通過除上 squre_root(d_k) 以歸一化點乘結果。

以下函數定義了一個標準的點乘注意力，該函數最終會返回匹配 Query 和 Key 的權重或概率 p_attn，以及最終注意力機制的輸出序列。

def attention(query, key, value, mask=None, dropout=None): "Compute Scaled Dot Product Attention" d_k = query.size(-1) scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) p_attn = F.softmax(scores, dim = -1) if dropout is not None: p_attn = dropout(p_attn) return torch.matmul(p_attn, value), p_attn

在上述函數中，query 矩陣的列數即維度數 d_k。在計算點乘並縮放後，我們可以在最後一個維度執行 Softmax 函數以得到概率 p_attn。

兩個最常見的注意力函數是加性注意力（additive attention）和點乘（乘法）注意力。除了要除上縮放因子 squre_root(d_k)，標準的點乘注意力與原論文中所採用的是相同的。加性注意力會使用單隱藏層的前饋網路計算適應度函數，它們在理論複雜度上是相似的。點積注意力在實踐中更快速且參數空間更高效，因為它能通過高度優化的矩陣乘法庫並行地計算。

Multi-head Attention

下圖展示了 Transformer 中所採用的 Multi-head Attention 結構，它其實就是多個點乘注意力並行地處理並最後將結果拼接在一起。一般而言，我們可以對三個輸入矩陣 Q、V、K 分別進行 h 個不同的線性變換，然後分別將它們投入 h 個點乘注意力函數並拼接所有的輸出結果。

Image(filename=images/ModalNet-20.png)

Multi-head Attention 允許模型聯合關注不同位置的不同表徵子空間信息，我們可以理解為在參數不共享的情況下，多次執行點乘注意力。Multi-head Attention 的表達如下所示：

其中 W 為對應線性變換的權重矩陣，Attention() 就是上文所實現的點乘注意力函數。

在原論文和實現中，研究者使用了 h=8 個並行點乘注意力層而完成 Multi-head Attention。對於每一個注意力層，原論文使用的維度是 d_k=d_v=d_model/h=64。由於每一個並行注意力層的維度降低，總的計算成本和單個點乘注意力在全維度上的成本非常相近。

以下定義了 Multi-head 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 # We assume d_v always equals d_k 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): "Implements Figure 2" if mask is not None: # Same mask applied to all h heads. mask = mask.unsqueeze(1) nbatches = query.size(0) # 1) Do all the linear projections in batch from d_model => h x d_k 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) Apply attention on all the projected vectors in batch. x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout) # 3) "Concat" using a view and apply a final linear. x = x.transpose(1, 2).contiguous() .view(nbatches, -1, self.h * self.d_k) return self.linears[-1](x)

在以上代碼中，首先我們會取 query 的第一個維度作為批量樣本數，然後再實現多個線性變換將 d_model 維的詞嵌入向量壓縮到 d_k 維的隱藏向量，變換後的矩陣將作為點乘注意力的輸入。點乘注意力輸出的矩陣將在最後一個維度拼接，即 8 個 n×64 維的矩陣拼接為 n×512 維的大矩陣，其中 n 為批量數。這樣我們就將輸出向量恢復為與詞嵌入向量相等的維度。

前面我們已經了解到 Transformer 使用了大量的自注意力機制，即 Attention(X, X, X )。簡單而言，Transformer 使用自注意力代替 RNN 或 CNN 抽取序列特徵。對於機器翻譯任務而言，自注意力輸入的 Query、Key 和 Value 都是相同的矩陣，那麼 Query 和 Key 之間的運算就相當於計算輸入序列內部的相似性，並根據這種相似性或權重注意到序列自身（Value）的內部聯繫。

這種內部聯繫可能是主語注意到謂語和賓語的信息或其它隱藏在句子內部的結構。Transformer 在神經機器翻譯和閱讀理解等任務上的優秀性能，都證明序列內部結構的重要性。

Transformer 以三種不同的方式使用 multi-head Attention。首先在編碼器到解碼器的層級中，Query 來源於前面解碼器的輸出，而記憶的 Key 與 Value 都來自編碼器的輸出。這允許解碼器中的每一個位置都注意輸入序列中的所有位置，因此它實際上模仿了序列到序列模型中典型的編碼器-解碼器注意力機制。

其次，編碼器包含了自注意力層，且該層中的所有 Value、Key 和 Query 都是相同的輸入矩陣，即編碼器的前層輸出。最後，解碼器中的自注意力層允許解碼器中的每一個位置都注意到包括當前位置的所有合法位置。這可以通過上文定義的 Mask 函數實現，從而防止產生左向信息流來保持自回歸屬性。

逐位置的前饋網路

為了注意子層，每一個編碼器和解碼器模塊最後都包含一個全連接前饋網路，它獨立且相同地應用於每一個位置。這個前饋網路包含兩個線性變換和一個非線性激活函數，且在訓練過程中我們可以在兩層網路之間添加 Dropout 方法：

如果我們將這兩個全連接層級與殘差連接和層級歸一化結合，那麼它就是每一個編碼器與解碼器模塊最後所必須的子層。我們可以將這一子層表示為：LayerNorm(x + max(0, x*w1 + b1)w2 + b2)。

儘管線性變換在所有不同的位置上都相同，但在不同的層級中使用不同的參數，這種變換其實同樣可以描述為核大小為 1 的兩個卷積。輸入和輸出的維度 d_model=512，而內部層級的維度 d_ff=2018。

如下所示，前饋網路的定義和常規的方法並沒有什麼區別，不過這個網路沒有添加偏置項，且對第一個全連接的輸出實現了 Dropout 以防止過擬合。

class PositionwiseFeedForward(nn.Module): "Implements FFN equation." 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))))

詞嵌入和 Softmax

與其它序列模型相似，我們可以使用學得的詞嵌入將輸入和輸出的辭彙轉換為維度等於 d_model 的向量。我們還可以使用一般的線性變換和 Softmax 函數將解碼器的輸出轉化為預測下一個辭彙的概率。在願論文的模型中，兩個嵌入層和 pre-softmax 線性變換的權重矩陣是共享的。在詞嵌入層中，我們將所有權重都乘以 squre_root(d_model)。

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 def forward(self, x): return self.lut(x) * math.sqrt(self.d_model)

位置編碼

位置編碼是 Transformer 模型中最後一個需要注意的結構，它對使用注意力機制實現序列任務也是非常重要的部分。如上文所述，Transformer 使用自注意力機制抽取序列的內部特徵，但這種代替 RNN 或 CNN 抽取特徵的方法有很大的局限性，即它不能捕捉序列的順序。這樣的模型即使能根據語境翻譯出每一個詞的意義，那也組不成完整的語句。

為了令模型能利用序列的順序信息，我們必須植入一些關於辭彙在序列中相對或絕對位置的信息。直觀來說，如果語句中每一個詞都有特定的位置，那麼每一個詞都可以使用向量編碼位置信息。將這樣的位置向量與詞嵌入向量相結合，那麼我們就為每一個詞引入了一定的位置信息，注意力機制也就能分辨出不同位置的詞。

谷歌研究者將「位置編碼」添加到輸入詞嵌入中，位置編碼有和詞嵌入相同的維度 d_model，每一個詞的位置編碼與詞嵌入向量相加可得出這個詞的最終編碼。目前有很多種位置編碼，包括通過學習和固定表達式構建的。

在這一項實驗中，谷歌研究者使用不同頻率的正弦和預先函數：

其中 pos 為詞的位置，i 為位置編碼向量的第 i 個元素。給定詞的位置 pos，我們可以將詞映射到 d_model 維的位置向量，該向量第 i 個元素就由上面兩個式子計算得出。也就是說，位置編碼的每一個維度對應於正弦曲線，波長構成了從 2π到 10000?2π的等比數列。

上面構建了絕對位置的位置向量，但詞的相對位置同樣非常重要，這也就是谷歌研究者採用三角函數表徵位置的精妙之處。正弦與餘弦函數允許模型學習相對位置，這主要根據兩個變換：sin(α+β)=sinα cosβ+cosα sinβ 以及 cos(α+β)=cosα cosβ?sinα sinβ。

對於辭彙間固定的偏移量 k，位置向量 PE(pos+k) 可以通過 PE(pos) 與 PE(k) 的組合表示，這也就表示了語言間的相對位置。

以下定義了位置編碼，其中我們對詞嵌入與位置編碼向量的和使用 Dropout，默認可令_drop=0.1。div_term 實現的是分母，而 pe[:, 0::2] 表示第二個維度從 0 開始以間隔為 2 取值，即偶數。

class PositionalEncoding(nn.Module): "Implement the PE function." def __init__(self, d_model, dropout, max_len=5000): super(PositionalEncoding, self).__init__() self.dropout = nn.Dropout(p=dropout) # Compute the positional encodings once in log space. 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) self.register_buffer(pe, pe) def forward(self, x): x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False) return self.dropout(x)

以下將基於一個位置將不同的正弦曲線添加到位置編碼向量中，曲線的頻率和偏移量在每個維度上都不同。

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]])None

谷歌等研究者在原論文中表示他們同樣對基於學習的位置編碼進行了實驗，並發現這兩種方法會產生幾乎相等的結果。所以按照性價比，他們還是選擇了正弦曲線，因為它允許模型在訓練中推斷更長的序列。

模型整體

下面，我們定義了一個函數以構建模型的整個過程，其中 make_model 在輸入原語辭彙表和目標語辭彙表後會構建兩個詞嵌入矩陣，而其它參數則會構建整個模型的架構。

def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1): "Helper: Construct a model from hyperparameters." c = copy.deepcopy attn = MultiHeadedAttention(h, d_model) ff = PositionwiseFeedForward(d_model, d_ff, dropout) position = PositionalEncoding(d_model, dropout) model = EncoderDecoder( Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N), Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N), nn.Sequential(Embeddings(d_model, src_vocab), c(position)), nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)), Generator(d_model, tgt_vocab)) # This was important from their code. # Initialize parameters with Glorot / fan_avg. for p in model.parameters(): if p.dim() > 1: nn.init.xavier_uniform(p) return model

在以上的代碼中，make_model 函數將調用上面我們定義的各個模塊，並將它們組合在一起。我們會將 Multi-Head Attention 子層、全連接子層和位置編碼等結構傳入編碼器與解碼器主體函數，再根據詞嵌入向量與位置編碼向量完成輸入與標註輸出的構建。以下簡單地示例了如何使用 make_model 函數構建模型：

# Small example model.tmp_model = make_model(10, 10, 2)None

訓練

這一部分將描述模型的訓練方案。首先需要介紹一些訓練標準編碼器解碼器模型的工具，例如定義一個批量的目標以儲存原語序列與目標語序列，並進行訓練。前文的模型架構與函數定義我們主要參考的原論文，而後面的具體訓練過程則主要參考了 Alexander 的實現經驗。

批量和掩碼

以下定義了保留一個批量數據的類，並且它會使用 Mask 在訓練過程中限制目標語的訪問序列。

class Batch: "Object for holding a batch of data with mask during training." def __init__(self, src, trg=None, pad=0): self.src = src self.src_mask = (src != pad).unsqueeze(-2) if trg is not None: self.trg = trg[:, :-1] self.trg_y = trg[:, 1:] self.trg_mask = self.make_std_mask(self.trg, pad) self.ntokens = (self.trg_y != pad).data.sum() @staticmethod def make_std_mask(tgt, pad): "Create a mask to hide padding and future words." tgt_mask = (tgt != pad).unsqueeze(-2) tgt_mask = tgt_mask & Variable( subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)) return tgt_mask

我們下一步需要創建一般的訓練和評分函數，以持續追蹤損失的變化。在構建一般的損失函數後，我們就能根據它更新參數。

如下定義了訓練中的迭代循環，我們使用 loss_compute() 函數計算損失函數，並且每運行 50 次迭代就輸出一次訓練損失，這有利於監控訓練情況。

def run_epoch(data_iter, model, loss_compute): "Standard Training and Logging Function" start = time.time() total_tokens = 0 total_loss = 0 tokens = 0 for i, batch in enumerate(data_iter): out = model.forward(batch.src, batch.trg, batch.src_mask, batch.trg_mask) loss = loss_compute(out, batch.trg_y, batch.ntokens) total_loss += loss total_tokens += batch.ntokens tokens += batch.ntokens if i % 50 == 1: elapsed = time.time() - start print("Epoch Step: %d Loss: %f Tokens per Sec: %f" % (i, loss / batch.ntokens, tokens / elapsed)) start = time.time() tokens = 0 return total_loss / total_tokens

訓練數據與分批

Alexander 等人的模型在標準的 WMT 2014 英語-德語數據集上進行訓練，這個數據集包含 450 萬條語句對。語句已經使用雙位元組編碼（byte-pair encoding）處理，且擁有約為 37000 個符號的原語-目標語共享辭彙庫。對於英語-法語的翻譯任務，. 原論文作者使用了更大的 WMT 2014 英語-法語數據集，它包含 3600 萬條語句，且將符號分割為包含 32000 個 word-piece 的辭彙庫。

原論文表示所有語句對將一同執行分批操作，並逼近序列長度。每一個訓練批量包含一組語句對，大約分別有 25000 個原語辭彙和目標語辭彙。

Alexander 等人使用 torch text 進行分批，具體的細節將在後面討論。下面的函數使用 torchtext 函數創建批量數據，並確保批量大小會填充到最大且不會超過閾值（使用 8 塊 GPU，閾值為 25000）。

global max_src_in_batch, max_tgt_in_batchdef batch_size_fn(new, count, sofar): "Keep augmenting batch and calculate total number of tokens + padding." global max_src_in_batch, max_tgt_in_batch if count == 1: max_src_in_batch = 0 max_tgt_in_batch = 0 max_src_in_batch = max(max_src_in_batch, len(new.src)) max_tgt_in_batch = max(max_tgt_in_batch, len(new.trg) + 2) src_elements = count * max_src_in_batch tgt_elements = count * max_tgt_in_batch return max(src_elements, tgt_elements)

batch_size_fn 將抽取批量數據，且每一個批量都抽取最大原語序列長度和最大目標語序列長度，如果長度不夠就使用零填充增加。

硬體與策略

原論文在一台機器上使用 8 塊 NVIDIA P100 GPU 訓練模型，基本模型使用了論文中描述的超參數，每一次迭代大概需要 0.4 秒。基本模型最後迭代了 100000 次，共花了 12 個小時。而對於大模型，每一次迭代需要花 1 秒鐘，所以訓練 300000 個迭代大概需要三天半。但我們後面的真實案例並不需要使用如此大的計算力，因為我們的數據集相對要小一些。

優化器

原論文使用了 Adam 優化器，其中β_1=0.9、 β_2=0.98 和 ?=10^{?9}。在訓練中，研究者會改變學習率為 l_rate=d?0.5model?min(step_num?0.5,step_num?warmup_steps?1.5)。

學習率的這種變化對應於在預熱訓練中線性地增加學習率，然後再與迭代數的平方根成比例地減小。這種 1cycle 學習策略在實踐中有非常好的效果，一般使用這種策略的模型要比傳統的方法收斂更快。在這個實驗中，模型採用的預熱迭代數為 4000。注意，這一部分非常重要，我們需要以以下配置訓練模型。

class NoamOpt: "Optim wrapper that implements rate." def __init__(self, model_size, factor, warmup, optimizer): self.optimizer = optimizer self._step = 0 self.warmup = warmup self.factor = factor self.model_size = model_size self._rate = 0 def step(self): "Update parameters and rate" self._step += 1 rate = self.rate() for p in self.optimizer.param_groups: p[lr] = rate self._rate = rate self.optimizer.step() def rate(self, step = None): "Implement `lrate` above" if step is None: step = self._step return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5)))def get_std_opt(model): return NoamOpt(model.src_embed[0].d_model, 2, 4000, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))

使用不同模型大小和最優化超參數下的變化曲線：

# Three settings of the lrate hyperparameters.opts = [NoamOpt(512, 1, 4000, None), NoamOpt(512, 1, 8000, None), NoamOpt(256, 1, 4000, None)]plt.plot(np.arange(1, 20000), [[opt.rate(i) for opt in opts] for i in range(1, 20000)])plt.legend(["512:4000", "512:8000", "256:4000"])None

正則化

• 標籤平滑

在訓練中，Alexander 等人使用了標籤平滑的方法，且平滑值?_ls=0.1。這可能會有損困惑度，因為模型將變得更加不確定它所做的預測，不過這樣還是提升了準確度和 BLEU 分數。

Harvard NLP 最終使用 KL 散度實現了標籤平滑，與其使用 one-hot 目標分布，他們選擇了創建一個對正確詞有置信度的分布，而其它平滑的概率質量分布將貫穿整個辭彙庫。

class LabelSmoothing(nn.Module): "Implement label smoothing." def __init__(self, size, padding_idx, smoothing=0.0): super(LabelSmoothing, self).__init__() self.criterion = nn.KLDivLoss(size_average=False) self.padding_idx = padding_idx self.confidence = 1.0 - smoothing self.smoothing = smoothing self.size = size self.true_dist = None def forward(self, x, target): assert x.size(1) == self.size true_dist = x.data.clone() true_dist.fill_(self.smoothing / (self.size - 2)) true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence) true_dist[:, self.padding_idx] = 0 mask = torch.nonzero(target.data == self.padding_idx) if mask.dim() > 0: true_dist.index_fill_(0, mask.squeeze(), 0.0) self.true_dist = true_dist return self.criterion(x, Variable(true_dist, requires_grad=False))

下面，我們可以了解到概率質量如何基於置信度分配到詞：

# Example of label smoothing.crit = LabelSmoothing(5, 0, 0.4)predict = torch.FloatTensor([[0, 0.2, 0.7, 0.1, 0], [0, 0.2, 0.7, 0.1, 0], [0, 0.2, 0.7, 0.1, 0]])v = crit(Variable(predict.log()), Variable(torch.LongTensor([2, 1, 0])))# Show the target distributions expected by the system.plt.imshow(crit.true_dist)None

標籤平滑實際上在模型對某些選項非常有信心的時候會懲罰它。

crit = LabelSmoothing(5, 0, 0.1)def loss(x): d = x + 3 * 1 predict = torch.FloatTensor([[0, x / d, 1 / d, 1 / d, 1 / d], ]) #print(predict) return crit(Variable(predict.log()), Variable(torch.LongTensor([1]))).data[0]plt.plot(np.arange(1, 100), [loss(x) for x in range(1, 100)])None

簡單的序列翻譯案例

我們可以從簡單的複製任務開始嘗試。若從小辭彙庫給定輸入符號的一個隨機集合，我們的目標是反向生成這些相同的符號。

def data_gen(V, batch, nbatches): "Generate random data for a src-tgt copy task." for i in range(nbatches): data = torch.from_numpy(np.random.randint(1, V, size=(batch, 10))) data[:, 0] = 1 src = Variable(data, requires_grad=False) tgt = Variable(data, requires_grad=False) yield Batch(src, tgt, 0)

計算模型損失

class SimpleLossCompute: "A simple loss compute and train function." def __init__(self, generator, criterion, opt=None): self.generator = generator self.criterion = criterion self.opt = opt def __call__(self, x, y, norm): x = self.generator(x) loss = self.criterion(x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1)) / norm loss.backward() if self.opt is not None: self.opt.step() self.opt.optimizer.zero_grad() return loss.data[0] * norm

貪婪解碼

# Train the simple copy task.V = 11criterion = LabelSmoothing(size=V, padding_idx=0, smoothing=0.0)model = make_model(V, V, N=2)model_opt = NoamOpt(model.src_embed[0].d_model, 1, 400, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))for epoch in range(10): model.train() run_epoch(data_gen(V, 30, 20), model, SimpleLossCompute(model.generator, criterion, model_opt)) model.eval() print(run_epoch(data_gen(V, 30, 5), model, SimpleLossCompute(model.generator, criterion, None)))Epoch Step: 1 Loss: 3.023465 Tokens per Sec: 403.074173Epoch Step: 1 Loss: 1.920030 Tokens per Sec: 641.6893801.9274832487106324Epoch Step: 1 Loss: 1.940011 Tokens per Sec: 432.003378Epoch Step: 1 Loss: 1.699767 Tokens per Sec: 641.9796651.657595729827881Epoch Step: 1 Loss: 1.860276 Tokens per Sec: 433.320240Epoch Step: 1 Loss: 1.546011 Tokens per Sec: 640.5371981.4888023376464843Epoch Step: 1 Loss: 1.278768 Tokens per Sec: 433.568756Epoch Step: 1 Loss: 1.062384 Tokens per Sec: 642.5420670.9853351473808288Epoch Step: 1 Loss: 1.269471 Tokens per Sec: 433.388727Epoch Step: 1 Loss: 0.590709 Tokens per Sec: 642.8621350.34273059368133546

這些代碼將簡單地使用貪婪解碼預測譯文。

def greedy_decode(model, src, src_mask, max_len, start_symbol): memory = model.encode(src, src_mask) ys = torch.ones(1, 1).fill_(start_symbol).type_as(src.data) for i in range(max_len-1): out = model.decode(memory, src_mask, Variable(ys), Variable(subsequent_mask(ys.size(1)) .type_as(src.data))) prob = model.generator(out[:, -1]) _, next_word = torch.max(prob, dim = 1) next_word = next_word.data[0] ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1) return ysmodel.eval()src = Variable(torch.LongTensor([[1,2,3,4,5,6,7,8,9,10]]) )src_mask = Variable(torch.ones(1, 1, 10) )print(greedy_decode(model, src, src_mask, max_len=10, start_symbol=1)) 1 2 3 4 5 6 7 8 9 10[torch.LongTensor of size 1x10]

真實案例

現在，我們將使用 IWSLT 德語-英語數據集實現翻譯任務。該任務要比論文中討論的 WMT 任務稍微小一點，但足夠展示整個系統。我們同樣還展示了如何使用多 GPU 處理來令加速訓練過程。

#!pip install torchtext spacy#!python -m spacy download en#!python -m spacy download de

數據載入

我們將使用 torchtext 和 spacy 載入數據集，並實現分詞。

# For data loading.from torchtext import data, datasetsif True: import spacy spacy_de = spacy.load(de) spacy_en = spacy.load(en) def tokenize_de(text): return [tok.text for tok in spacy_de.tokenizer(text)] def tokenize_en(text): return [tok.text for tok in spacy_en.tokenizer(text)] BOS_WORD = <s> EOS_WORD = </s> BLANK_WORD = "<blank>" SRC = data.Field(tokenize=tokenize_de, pad_token=BLANK_WORD) TGT = data.Field(tokenize=tokenize_en, init_token = BOS_WORD, eos_token = EOS_WORD, pad_token=BLANK_WORD) MAX_LEN = 100 train, val, test = datasets.IWSLT.splits( exts=(.de, .en), fields=(SRC, TGT), filter_pred=lambda x: len(vars(x)[src]) <= MAX_LEN and len(vars(x)[trg]) <= MAX_LEN) MIN_FREQ = 2 SRC.build_vocab(train.src, min_freq=MIN_FREQ) TGT.build_vocab(train.trg, min_freq=MIN_FREQ)

我們希望有非常均勻的批量，且有最小的填充，因此我們必須對默認的 torchtext 分批函數進行修改。這段代碼修改了默認的分批過程，以確保我們能搜索足夠的語句以找到緊湊的批量。

數據迭代器

迭代器定義了分批過程的多項操作，包括數據清洗、整理和分批等。

class MyIterator(data.Iterator): def create_batches(self): if self.train: def pool(d, random_shuffler): for p in data.batch(d, self.batch_size * 100): p_batch = data.batch( sorted(p, key=self.sort_key), self.batch_size, self.batch_size_fn) for b in random_shuffler(list(p_batch)): yield b self.batches = pool(self.data(), self.random_shuffler) else: self.batches = [] for b in data.batch(self.data(), self.batch_size, self.batch_size_fn): self.batches.append(sorted(b, key=self.sort_key))def rebatch(pad_idx, batch): "Fix order in torchtext to match ours" src, trg = batch.src.transpose(0, 1), batch.trg.transpose(0, 1) return Batch(src, trg, pad_idx)

下面，我們利用前面定義的函數創建了模型、度量標準、優化器、數據迭代器和並行化：

# GPUs to usedevices = [0, 1, 2, 3]if True: pad_idx = TGT.vocab.stoi["<blank>"] model = make_model(len(SRC.vocab), len(TGT.vocab), N=6) model.cuda() criterion = LabelSmoothing(size=len(TGT.vocab), padding_idx=pad_idx, smoothing=0.1) criterion.cuda() BATCH_SIZE = 12000 train_iter = MyIterator(train, batch_size=BATCH_SIZE, device=0, repeat=False, sort_key=lambda x: (len(x.src), len(x.trg)), batch_size_fn=batch_size_fn, train=True) valid_iter = MyIterator(val, batch_size=BATCH_SIZE, device=0, repeat=False, sort_key=lambda x: (len(x.src), len(x.trg)), batch_size_fn=batch_size_fn, train=False) model_par = nn.DataParallel(model, device_ids=devices)None

下面可以訓練模型了，Harvard NLP 團隊首先運行了一些預熱迭代，但是其它的設定都能使用默認的參數。在帶有 4 塊 Tesla V100 的 AWS p3.8xlarge 中，批量大小為 12000 的情況下每秒能運行 27000 個詞。

訓練系統

#!wget https://s3.amazonaws.com/opennmt-models/iwslt.ptif False: model_opt = NoamOpt(model.src_embed[0].d_model, 1, 2000, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)) for epoch in range(10): model_par.train() run_epoch((rebatch(pad_idx, b) for b in train_iter), model_par, MultiGPULossCompute(model.generator, criterion, devices=devices, opt=model_opt)) model_par.eval() loss = run_epoch((rebatch(pad_idx, b) for b in valid_iter), model_par, MultiGPULossCompute(model.generator, criterion, devices=devices, opt=None)) print(loss)else: model = torch.load("iwslt.pt")

一旦訓練完成了，我們就能解碼模型並生成一組翻譯，下面我們簡單地翻譯了驗證集中的第一句話。該數據集非常小，所以模型通過貪婪搜索也能獲得不錯的翻譯效果。

for i, batch in enumerate(valid_iter): src = batch.src.transpose(0, 1)[:1] src_mask = (src != SRC.vocab.stoi["<blank>"]).unsqueeze(-2) out = greedy_decode(model, src, src_mask, max_len=60, start_symbol=TGT.vocab.stoi["<s>"]) print("Translation:", end=" ") for i in range(1, out.size(1)): sym = TGT.vocab.itos[out[0, i]] if sym == "</s>": break print(sym, end =" ") print() print("Target:", end=" ") for i in range(1, batch.trg.size(0)): sym = TGT.vocab.itos[batch.trg.data[i, 0]] if sym == "</s>": break print(sym, end =" ") print() breakTranslation: <unk> <unk> . In my language , that means , thank you very much . Gold: <unk> <unk> . It means in my language , thank you very much .

實驗結果

在 WMT 2014 英語到法語的翻譯任務中，原論文中的大型的 Transformer 模型實現了 41.0 的 BLEU 分值，它要比以前所有的單模型效果更好，且只有前面頂級的模型 1/4 的訓練成本。在 Harvard NLP 團隊的實現中，OpenNMT-py 版本的模型在 EN-DE WMT 數據集上實現了 26.9 的 BLEU 分值。