想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

pythontab發表於2019-04-04

目前,無論是從效能、結構還是業界應用上,Transformer 都有很多無可比擬的優勢。本文將介紹 PaddlePaddle 的 Transformer 專案,我們從專案使用到原始碼解析帶你玩一玩 NMT。只需千行模型程式碼,Transformer 實現帶回家。

其實 PyTorch、TensorFlow 等主流框架都有 Transformer 的實現,但如果我們需要將它們應用到產品中,還是需要修改很多。

例如谷歌大腦構建的 Tensor2Tensor,它最開始是為了實現 Transformer,後來擴充套件到了各種任務。對於基於 Tensor2Tensor 實現翻譯任務的使用者,他們需要在 10 萬+行 TensorFlow 程式碼找到需要的部分。

PaddlePaddle 提供的 Transformer 實現,專案程式碼只有 2000+行,簡潔優雅。如果我們使用大 Batch Size,那麼在預測速度上,PaddlePaddle 復現的模型比 TensorFlow 官方使用 Tensor2Tensor 實現的模型還要快 4 倍。

專案地址:https://github.com/PaddlePaddle/models/tree/develop/fluid/PaddleNLP/neural_machine_translation/transformer

1. Transformer 怎麼用

相比此前 Seq2Seq 模型中廣泛使用的迴圈神經網路,Transformer 使用深層注意力機制獲得了更好的效果,目前大多數神經機器翻譯模型都採用了這一網路結構。此外,不論是新興的預訓練語言模型,還是問答或句法分析,Transformer 都展現出強大的建模能力。

相比傳統 NMT 使用迴圈層或卷積層抽取文字資訊,Transformer 使用自注意力網路抽取並表徵這些資訊,下圖對比了不同層級的特點: 

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

不同網路的主要性質,其中 n 表示序列長度、d 為隱向量維度、k 為卷積核大小。例如單層計算複雜度,一般句子長度 n 都小於隱向量維度 d,那麼自注意力層級的計算複雜度最小。

如上所示,Transformer 使用的自注意力模型主要擁有以下優點,1)網路結構的計算複雜度最低;2)由於序列運算元複雜度低,模型的並行度很高;3)最大路徑長度小,能夠更好地表示長距離依賴關係;4)模型更容易訓練。

現在,如果我們需要訓練一個 Transformer,那麼最好的方法是什麼?當然是直接跑已復現的模型了,下面我們將跑一跑 PaddlePaddle 實現的 Transformer。

1.1 處理資料

在 PaddlePaddle 的復現中,百度採用原論文測試的 WMT'16 EN-DE 資料集,它是一箇中等規模的資料集。這裡比較方便的是,百度將資料下載和預處理等過程都放到了 gen_data.sh 指令碼中,包括 Tokenize 和 BPE 編碼。

在這個專案中,我們既可以通過指令碼預處理資料,也可以使用百度預處理好的資料集。首先最簡單的方式是直接執行 gen_data.sh 指令碼,執行後可以生成 gen_data 資料夾,該資料夾主要包含以下檔案:

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

其中 wmt16_ende_data_bpe 資料夾包含最終使用的英德翻譯資料。

如果我們從頭下載並預處理資料,那麼大概需要花 1 到 2 個小時完成預處理。為此,百度也提供了預處理好的 WMT'16 EN-DE 資料集,它包含訓練、驗證和測試所需要的 BPE 資料和字典。

其中,BPE 策略會把稀疏詞拆分為高頻的子詞,這樣既能解決低頻詞無法訓練的問題,也能合理降低詞表規模。

如果不採用 BPE 的策略,要麼詞表的規模變得很大,從而使訓練速度變慢或者視訊記憶體太小而無法訓練;要麼一些低頻詞會當作未登入詞處理,從而得不到訓練。

預處理資料地址:https://transformer-res.bj.bcebos.com/wmt16_ende_data_bpe_clean.tar.gz

如果我們有其它資料集,例如中英翻譯資料,也可以根據特定的格式進行定義。例如用空格分隔不同的 token(對於中文而言需要提前用分詞工具進行分詞),用\t 分隔源語言與目標語句對。

1.2 訓練模型

如果需要執行模型訓練,我們也可以直接執行訓練主函式 train.py。如下簡要配置了資料路徑以及各種模型引數

# 視訊記憶體使用的比例,視訊記憶體不足可適當增大,最大為1
export FLAGS_fraction_of_gpu_memory_to_use=0.8
# 視訊記憶體清理的閾值,視訊記憶體不足可適當減小,最小為0,為負數時不啟用
export FLAGS_eager_delete_tensor_gb=0.7
python -u train.py \
  --src_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \
  --trg_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \
  --special_token '<s>' '<e>' '<unk>' \
  --train_file_pattern gen_data/wmt16_ende_data_bpe/train.tok.clean.bpe.32000.en-de \
  --token_delimiter ' ' \
  --use_token_batch True \
  --batch_size 1600 \
  --sort_type pool \
  --pool_size 200000 \
  n_head 8 \
  n_layer 4 \
  d_model 512 \
  d_inner_hid 1024 \
  prepostprocess_dropout 0.3

此外,如果視訊記憶體不夠大,那麼我們可以將 Batch Size 減小一點。為了快速測試訓練效果,我們將模型調得比 Base Transformer 還小(降低網路的層數、head 的數量、以及隱層的大小)。

上面僅展示了小部分的超參設定,更多的配置可以在 GitHub 專案 config.py 檔案中找到。預設情況下,模型每迭代一萬次儲存一次模型,每個 epoch 結束後也會儲存一次 cheekpoint。此外,在我們訓練的過程中,預設每一百次迭代會列印一次模型資訊,其中 ppl 表示的是困惑度,困惑度越小模型效果越好。

在單機訓練中,預設使用所有 GPU,可以通過 CUDA_VISIBLE_DEVICES 環境變數來設定使用的 GPU,例如 CUDA_VISIBLE_DEVICES='0,1',表示使用 0 號和 1 號卡進行訓練。

1.3 預測推斷

訓練完 Transformer 後就可以執行推斷了,我們需要執行對應的推斷檔案 infer.py。我們也可以在推斷過程中配置超引數,但注意超參需要和前面訓練時保持一致。

python -u infer.py \
  --src_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \
  --trg_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \
  --special_token '<s>' '<e>' '<unk>' \
  --test_file_pattern gen_data/wmt16_ende_data_bpe/newstest2016.tok.bpe.32000.en-de \
  --token_delimiter ' ' \
  --batch_size 32 \
  model_path trained_models/iter_100000.infer.model \
  n_head 8 \
  n_layer 4 \
  d_model 512 \
  d_inner_hid 1024 \
  prepostprocess_dropout 0.3
  beam_size 5 \
  max_out_len 255

相比模型的訓練,推斷過程需要一些額外的超引數,例如配置 model_path 指定模型所在目錄、設定 beam_size 和 max_out_len 來指定 Beam Search 每一步候選詞的個數和最大翻譯長度。這些超引數也可以在 config.py 中找到,該檔案對這些超參都有註釋說明。

執行以上預測命令會將翻譯結果直接打出來,每行輸出是對應行輸入得分最高的翻譯。對於使用 BPE 的英德資料,預測出的翻譯結果也將是 BPE 表示的資料,所以需要還原成原始資料才能進行正確評估。如下命令可以將 predict.txt 內的翻譯結果(BPE 表示)恢復到 predict.tok.txt 檔案中(tokenize 後的資料):

sed -r 's/(@@ )|(@@ ?$)//g' predict.txt > predict.tok.txt

在未使用整合方法的情況下,百度表示 base model 和 big model 在收斂後,測試集的 BLEU 值參考如下:

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

這兩個預訓練模型也提供了下載地址:

  • Base:https://transformer-res.bj.bcebos.com/base_model.tar.gz

  • Big:https://transformer-res.bj.bcebos.com/big_model.tar.gz

2. Transformer 怎麼改

如果我們想要訓練自己的 Transformer,那麼又該怎樣理解並修改 PaddlePaddle 程式碼呢?如果我們需要根據自己的資料集和任務改程式碼,除了前面資料預處理過程,模型結構等模組有時也需要修改。這就需要我們先理解原始碼了,PaddlePaddle 的原始碼基本都是基礎的函式或運算,我們很容易理解並使用。

對於 PaddlePaddle 不熟悉的讀者可查閱文件,也可以看看入門教程,瞭解基本編寫模式後就可以看懂整個實現了。

PaddlePaddle 官網地址:http://paddlepaddle.org/paddle

如 Seq2Seq 一樣,原版 Transformer 也採用了編碼器-解碼器框架,但它們會使用多個 Multi-Head 注意力、前饋網路、層級歸一化和殘差連線等。下圖從左到右展示了原論文所提出的 Transformer 架構、Multi-Head 注意力和標量點乘注意力。

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

上圖右邊的點乘注意力就是標準 Seq2Seq 模型中的注意力機制,中間的 Multi-head 注意力其實就是將一個隱層資訊切分為多份,並單獨計算注意力資訊,使得一個詞與其它多個目標詞的注意力資訊計算更精確。最左邊為 Transformer 的整體架構,編碼器與解碼器由多個類似的模組組成,後面將簡要介紹這些模組與對應的 PaddlePaddle 程式碼。

2.1 點乘注意力

注意力機制目前在機器翻譯中已經極其流行了,我們可以認為 Transformer 是一種堆疊多層注意力網路的模型,它採用的是一種名為經縮放的點乘注意力機制。這種注意力機制使用經縮放的點乘作為作為評分函式,從而評估各隱藏狀態對當前預測的重要性,如下是該注意力的表示式:

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

其中 Query 向量與 (Key, Value ) 向量在 NMT 中相當於目標語輸入序列與源語輸入序列,Query 與 Key 向量的點乘,經過 SoftMax 函式後可得出一組歸一化的概率。這些概率相當於給源語輸入序列做加權平均,即表示在生成新的隱層資訊的時候需要關注哪些詞。

在 Transformer 的 PaddlePaddle 實現中,經縮放的點乘注意力是在 Multi-head 注意力函式下實現的,如下所示為上述表示式的實現程式碼:

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

在這個函式中,q、k、v 和公式中的一樣,attn_bias 用於 Mask 掉選定的特定位置(encode 的 self attention 和 decoder 端的 encode attention 都是遮蔽掉 padding 的詞;decoder 的 self attention 遮蔽掉當前詞後面的詞,目的是為了和解碼的過程保持一致),因此在給不同輸入加權時忽略該位置的輸入。

如上 product 計算的是 q 和 k 之間的點乘,且經過根號下 d_key(key 的維度)的縮放。這裡我們可以發現引數 alpha 可以直接對矩陣乘法的結果進行縮放,預設情況下它為 1.0,即不進行縮放。在 Transformer 原論文中,作者表示如果 d_key 比較小,那麼直接點乘和帶縮放的點乘差別不大,所以他們認為高維情況下可能不帶縮放的乘積太大而令 Softmax 函式飽和。

weights 表示對輸入的不同元素加權,即不同輸入對當前預測的重要性,訓練中也可以對該權重進行 Dropout。最後 out 表示按照 weights 對輸入 V 進行加權和,得出來就是當前注意力的運算結果。

2.2 Muti-head 注意力

Multi-head 注意力其實就是多個點乘注意力並行地處理並最後將結果拼接在一起。一般而言,我們可以對三個輸入矩陣 Q、V、K 分別進行線性變換,然後分別將它們投入 h 個點乘注意力函式並拼接所有的輸出結果。

這種注意力允許模型聯合關注不同位置的不同表徵子空間資訊,我們可以理解為在引數不共享的情況下,多次執行點乘注意力。如下所示為 Muti-head 注意力的表示式:

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

其中每一個 head 都為一個點乘注意力,不同 head 的輸入是相同 Q、K、V 的不同線性變換。

總體而言,PaddlePaddle 的 Multi-head 注意力實現分為幾個步驟:先為 Q、K、V 執行線性變換;再變換維度以計算點乘注意力;最後計算各 head 的注意力輸出併合並在一起。

2.2.1 線性變換

如前公式所示,Muti-head 首先要執行線性變換,從而令不同的 head 關注不同表徵空間的資訊。這種線性變換即乘上不同的權重矩陣,且模型在訓練過程中可以學習和更新這些權重矩陣。在如下的 PaddlePaddle 程式碼中,我們可以直接呼叫全連線層 layers.fc() 完成線性變換。

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

直接呼叫全連線層會自動為輸入建立權重,且我們要求不使用偏置項和啟用函式。這裡比較方便的是,PaddlePaddle 的 layers.fc() 函式可以接受高維輸入,省略了手動展平輸入向量的操作。因此這裡有 num_flatten_dims=2,即將前兩個維度展平為一個維度,第三個維度保持不變。

例如對於輸入張量 q 而言,線性變換的輸出維度應該是 [batch_size,max_sequence_length,d_key * n_head],最後一個維度即 n_head 個 d_key 維的 Query 向量。每一個 d_key 維的向量都會饋送到不同的 head,並最後拼接起來。

2.2.2 維度變換

為了進行 Multi-Head 的運算,我們需要將線性變換的結果進行 reshape 和轉置操作。現在我們將這幾個張量的最後一個維度分割成不同的 head,並做轉置以便於後續運算。

具體而言,輸入張量 q、k 和 v 的維度資訊為 [bs, max_sequence_length, n_head * hidden_dim],我們希望把它們轉換為 [bs, n_head, max_sequence_length, hidden_dim]。

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

如上使用 layers.reshape() 和 layers.transpose() 函式完成分割與轉置。其中 layers.reshape() 在接收輸入張量後會按照形狀 [0, 0, n_head, d_key] 進行轉換,其中 0 表示從輸入張量對應維數複製出來。此外,因為 inplace 設定為 True,那麼 reshape 操作就不會進行資料的複製,從而提升運算效率。

後面的轉置就比較簡單了,只需要按照維度索引將第「1」個維度和第「2」個維度交換就行了。此外為了更快地執行推斷,PaddlePaddle 實現程式碼還做了非常多的優化,例如這部分後續會對推斷過程的快取和處理流程進行優化。

2.2.3 合併

前面已經介紹過點乘注意力了,那麼上面對 q、k、v 執行維度變換後就可直接傳入點乘注意力函式,並計算出 head_1、head_2 等注意力結果。現在最後一步只需要將這些 head 拼接起來就完成了整個過程,也就完成了上面 Multi-head 注意力的計算式。

因為每一個批量、head 和時間步都會計算得出一個注意力向量,因此總體上注意力計算結果的維度資訊為 [bs, n_head, max_sequence_length, hidden_dim]。如果要將不同的 head 拼接在一起,即將 head 這個維度合併到 hidden_dim 中去,因此合併的過程和前面維度變換的過程正好相反。

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

如上合併過程會先檢驗維度資訊,然後先轉置再 reshape 合併不同的 head。注意在原論文中,合併不同的 head 後,還需要再做一個線性變換,這個線性變換的結果就是 Muti-head 注意力的輸出了。

最後,我們再將上面的四部分串起來就是 Transformer 最核心的 Multi-head 注意力。理解了各個模組後,下面串起來就能愉快地看懂整個過程了:

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

當然,如果編碼器和解碼器輸入到 Multi-head 注意力的 q 與 (k、v) 是相同的,那麼它又可稱為自注意力網路。

2.3 前饋網路

對於每一個編碼器和解碼器模組,除了殘差連線與層級歸一化外,重要的就是堆疊 Muti-head 注意力和前饋網路(FFN)。前面我們已經解決了 Multi-head 注意力,現在需要理解主位置的前饋網路了。直觀而言,FFN 的作用是整合 Multi-head 注意力生成的上下文向量,因此能更好地利用從源語句子和目標語句子抽取的深度資訊。

如下所示在原論文中,前饋網路的計算過程可以表達為以下方程:

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

前饋網路的結構很簡單,一個 ReLU 啟用函式加兩次線性變換就完成了。如下基本上只需要呼叫 PaddlePaddle 的 layers.fc() 就可以了:

想要千行程式碼搞定Transformer?這份高效的PaddlePaddle官方實現請收下

現在基本上核心操作就定義完了,後面還有更多模組與架構,例如怎樣利用核心操作搭建編碼器模組與解碼器模組、如何搭建整體 Transformer 模型等,讀者可繼續閱讀原專案中的簡潔程式碼。整體而言,包括上面程式碼在內,千行程式碼就可以完全弄懂 Transformer,PaddlePaddle 的 Transformer 復現值得我們仔細讀一讀。

此外,在這千行模型程式碼中,為了給訓練和推斷加速,還有很多特殊技巧。例如在 Decoder 中加入對 Encoder 計算結果的快取等。加上這些技巧,PaddlePaddle 的實現才能在大 Batch Size 下實現 4 倍推斷加速。

因為本身 PaddlePaddle 程式碼就已經非常精煉,通過它們也很容易理解這些技巧。基本上看函式名稱就能知道大致的作用,再結合文件使用就能完全讀懂了。

最後,除了模型架構,整個專案還會有其它組成部分,例如訓練、推斷、資料預處理等等。這些程式碼同樣非常簡潔,我們可以根據實際需求閱讀並修改它們。總體而言,PaddlePaddle 的 Transformer 實現確實非常適合理解與修改。想要跑一跑神經機器翻譯的同學,PaddlePaddle 的 Transformer 實現確實值得推薦。

相關文章