用于NAT的選擇性知識(shí)蒸餾框架

01

研究動(dòng)機(jī)

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

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

圖1：Selective KD的流程示意圖

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

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

02

解決方案

2.1評(píng)估模型

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

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

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

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

嚴(yán)重的錯(cuò)誤：翻譯的質(zhì)量很糟糕

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

圖2：4種不同的情況對(duì)應(yīng)的案例

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

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

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

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

03

實(shí)驗(yàn)

我們?cè)赪MT14 EN-DE和WMT16 EN-RO上開展了實(shí)驗(yàn)，包括了兩種代表性的NAT架構(gòu)：CMLM [5]和GLAT+CTC [6]，以及一種inference-efficient的AT架構(gòu)：DeepShallow [7]（6層編碼器，1層解碼器）。

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

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

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

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

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

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

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

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

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

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

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

04

總結(jié)

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

審核編輯 ：李倩