用于NAT的選擇性知識蒸餾框架

01

研究動機

在本文中，我們研究了一種能夠高效推理的機器翻譯模型NAT (Non-Autoregressive Transformer)[1]。相較于傳統(tǒng)的Transformer，NAT能夠在解碼階段并行預測，從而大幅提升模型的推理速度。此外，NAT可以使得模型在訓練和測試階段從相同的分布進行預測，從而有效避免了順序解碼模型中經(jīng)常出現(xiàn)的exposure bias問題。在WMT21 news translation shared task for German→English translation中，已經(jīng)有NAT模型在翻譯質(zhì)量上超過了許多順序解碼的模型。

盡管NAT在擁有許多潛在的優(yōu)勢，目前的工作中這類模型仍然在很大程度上依賴于句子級別的知識蒸餾（sequence-level knowledge distillation, KD）[2]。由于需要并行預測所有token，NAT對單詞間依賴關系的建模能力較弱。這個特點使得在真實數(shù)據(jù)集上，NAT很容易受到multi-modality問題的影響：訓練數(shù)據(jù)中一個輸入可能對應多個不同的輸出。在這樣的背景下，Gu提出訓練一個AT (Autoregressive Transformer)[3]模型作為老師，將它的輸出作為NAT的學習對象。這種KD方式可以幫助NAT繞過multi-modality問題，從而大幅提升NAT的翻譯表現(xiàn)。

圖1：Selective KD的流程示意圖

KD在幫助NAT提升表現(xiàn)的同時，也會帶來一些負面影響，例如模型在低頻詞上的準確率較低[4]、AT teacher的錯誤會傳播到NAT上等。此外，如果NAT僅能在AT teacher的輸出上學習，這類模型的翻譯質(zhì)量將很難有更進一步的突破。我們的研究希望能夠在避免multi-modality的情況下，讓NAT能夠從真實的數(shù)據(jù)分布中學到知識蒸餾的過程中缺失的信息，從而提升NAT的表現(xiàn)。

為達到這樣的目的，我們提出了selective KD：在KD數(shù)據(jù)上訓練一個NAT作為評估模型，并通過它來選擇需要蒸餾的句子。通過這種方式，我們可以讓模型接觸到翻譯質(zhì)量更高的真實數(shù)據(jù)，同時避免了嚴重的multi-modality情況。受課程學習的影響，我們也在訓練過程中動態(tài)調(diào)整蒸餾數(shù)據(jù)的比例。“用評估模型有選擇地蒸餾數(shù)據(jù)”和“動態(tài)調(diào)節(jié)蒸餾數(shù)據(jù)的比例”共同構成了我們的Selective KD訓練框架。

02

解決方案

2.1評估模型

我們首先將數(shù)據(jù)蒸餾產(chǎn)生的結果劃分為四種不同的情況：

較輕的modality change：某些單詞可能被替換為同義詞，句式和語義并沒有發(fā)生顯著的變化

較輕的錯誤：在保持原有句式和語義的情況下，發(fā)生了一些小錯誤，例如單詞重復

嚴重的modality change：語義不變的情況下，句子的表達方式發(fā)生了顯著的變化

嚴重的錯誤：翻譯的質(zhì)量很糟糕

對于情況1，我們可以容忍較輕的modality change，這種情況下真實數(shù)據(jù)和蒸餾數(shù)據(jù)都可以被視作正確的學習目標，同時引入真實數(shù)據(jù)不會大幅增加數(shù)據(jù)集的復雜程度。情況2中，用真實數(shù)據(jù)替換蒸餾數(shù)據(jù)可以得到更高的翻譯質(zhì)量，找出屬于這種情況的樣本是我們方法的主要目標。情況3中，由于引入真實數(shù)據(jù)會惡化multi-modality問題，我們希望蒸餾這部分數(shù)據(jù)。情況4很少發(fā)生，我們認為這種情況下該訓練樣本對NAT可能太過困難，引入真實數(shù)據(jù)帶來的提升很有限?？偟膩碚f，我們希望能找到情況1、2對應的訓練樣本，在訓練過程中將它們的原始數(shù)據(jù)作為學習對象。

圖2：4種不同的情況對應的案例

為了篩選情況1、2中的數(shù)據(jù)，我們在蒸餾數(shù)據(jù)上訓練一個NAT作為評估模型，通過比較評估模型的輸出和真實數(shù)據(jù)計算一個score，判斷一個真實翻譯是否適合被直接用于訓練。若對于某個樣本評估模型的輸出和真實數(shù)據(jù)較為接近，則score較高，我們可以認為蒸餾數(shù)據(jù)僅有微小的錯誤或modality change，從而認為它屬于情況1、2，無需蒸餾。反之，可以認為蒸餾數(shù)據(jù)發(fā)生了較大的變化，因此屬于情況3、4，或是這個樣本在蒸餾后不發(fā)生太大變化的情況下對NAT而言仍過于困難。經(jīng)過篩選，我們僅蒸餾那些不適合用于訓練的真實數(shù)據(jù)。

2.2動態(tài)調(diào)整蒸餾比例：由困難到容易

我們在訓練過程中會調(diào)整蒸餾數(shù)據(jù)的比例。一般來說，剛開始訓練時絕大多數(shù)訓練樣本為真實數(shù)據(jù)，訓練的尾聲則會蒸餾整個訓練集。具體實現(xiàn)中，我們通過動態(tài)調(diào)節(jié)score的閾值來調(diào)整蒸餾的比例。

圖3：selective KD在第k次update的算法示意

03

實驗

我們在WMT14 EN-DE和WMT16 EN-RO上開展了實驗，包括了兩種代表性的NAT架構：CMLM [5]和GLAT+CTC [6]，以及一種inference-efficient的AT架構：DeepShallow [7]（6層編碼器，1層解碼器）。

3.1翻譯質(zhì)量與推理速度

我們通過BLEU score [8]和一種learned metric COMET [9]來衡量模型的翻譯質(zhì)量，并通過和標準Transformer比較來衡量推理速度。可以發(fā)現(xiàn)，相比于常規(guī)的知識蒸餾，Selective KD可以在不同數(shù)據(jù)集、不同架構以及不同metric上穩(wěn)定取得翻譯質(zhì)量的提升，同時保持模型自身在推理速度上的優(yōu)勢。我們方法在inference-efficient AT上也有明顯的效果，這進一步說明了selective KD具有廣泛的價值。

圖4：翻譯質(zhì)量與推理速度。翻譯質(zhì)量括號外為BLEU，括號內(nèi)為COMET

3.2調(diào)節(jié)quality和complexity

真實數(shù)據(jù)的翻譯質(zhì)量往往是優(yōu)于蒸餾數(shù)據(jù)的，通過調(diào)節(jié)蒸餾數(shù)據(jù)的比例，Selective KD可以調(diào)節(jié)訓練集的quality。與此同時，我們希望知道這個方法是否可以靈活調(diào)節(jié)訓練集的complexity。為了更好地觀察這一點，文章中用了兩個metric來衡量數(shù)據(jù)的復雜程度：Translatioin Uncertainty [10]和Alignment Shift。Translation Uncertainty反映了源句單詞對應翻譯結果的多樣性，Alignment Shift反映了句式的變化程度。

圖5：Translation Uncertainty（左）和Alignment Shift（右）的計算方式

如圖6所示，我們的方法可以有效控制數(shù)據(jù)的complexity。我們保留的真實數(shù)據(jù)（綠色折線）在兩個指標上都遠遠低于被蒸餾的真實數(shù)據(jù)（紅色折線）。在增加真實數(shù)據(jù)的比例同時，整個數(shù)據(jù)集complexity的提升是緩慢而平滑的。

圖6：數(shù)據(jù)的Translation Uncertainty（左）和Alignment Shift（右）

3.3蒸餾數(shù)據(jù)占比的影響

如圖7所示，我們在不同蒸餾比例的數(shù)據(jù)上進行了實驗?？梢园l(fā)現(xiàn)，通過selective KD僅蒸餾5%的數(shù)據(jù)就可以提升2.4 BLEU。在蒸餾數(shù)據(jù)比例為80%時，模型的表現(xiàn)甚至超過了完全蒸餾的數(shù)據(jù)，根據(jù)[10]，一種可能的解釋是這種比例下數(shù)據(jù)的complexity更適合我們實驗中采用的GLAT+CTC架構。另外，動態(tài)調(diào)節(jié)真實數(shù)據(jù)的比例（藍色虛線）可以進一步提升模型的表現(xiàn)。

圖7：在不同蒸餾比例下模型的表現(xiàn)

04

總結

在這篇文章中，我們提出了選擇性知識蒸餾，從而使得NAT模型可以從真實的數(shù)據(jù)分布中學到知識蒸餾過程中缺失的部分信息。具體來說，我們采用一個NAT作為評估模型來判斷哪些句子需要蒸餾，并動態(tài)提高蒸餾數(shù)據(jù)的比例。我們用實驗結果證明了該方法可以有效提升NAT在機器翻譯任務上的表現(xiàn)。

審核編輯 ：李倩