深度模型在神經機器翻譯（Neural Machine Translation，NMT）中取得突破性進展，但其計算量大且佔用大量記憶體。基於此，講者提出一種基於群體置換的知識蒸餾方法（Group-Permutation based Knowledge Distillation，GPKD）， 將深層的Transformer模型壓縮為淺層模型；為了進一步增強teacher模型的效能，提出了子層跳躍的方法（Skipping Sub-Layer），透過在三個資料集上的實驗，驗證了該兩種方法的有效性。

李北，東北大學自然語言處理實驗室博士二年級學生，導師是肖桐教授和朱靖波教授。研究方向主要包括深層網路建模，篇章級機器翻譯，模型結構最佳化等話題。目前以第一作者在ACL、EMNLP、AAAI等會議發表多篇論文。

一、研究背景與動機

在自然語言處理的各項任務中均取得突破性進展的Transformer模型，採用自注意力網路與前饋神經網路的結構，具有高效的平行計算能力。儘管Transformer十分強大，但應用場景與需求不同，可進一步完善。講者列舉了部分針對Transformer進行改進的工作：

（1）引入相對位置資訊，能夠有效彌補自注意力機制本身缺乏對相對資訊捕捉的能力，其中比較有代表性的工作是RPR以及Transformer-XL；

（2）從語言學角度來說，周圍詞與當前詞可能蘊含更直接的聯絡，基於此可採用卷積的思想進行自注意力機制的計算；

（3）在處理長文字時提高網路的計算效率，比如Linformer；

（4）更復雜的層級連線，也是近年受關注的一個研究方向。

圖1. 針對Transformer模型改進工作介紹

上述工作均是在受限模型理論下進行研究，然而神經網路本身是透過構造複雜函式來擬合數據，進而求解目標問題，因此模型容量對最終效能有很大的影響。這說明更大模型容量，能更好地利用資料來解決更復雜的問題。

常用的增大網路容量的方式大致可分為兩種：

第一種是增大模型的寬度，這也是早期Transformer-Big模型的一種設計的思路，利用更大的隱藏層維度和詞嵌入維度以及FFN中的一個對映維度來提高模型的表示能力；

第二種是增大模型的深度，其優勢在於透過更多的非線性與線性函式組合來構造更復雜的函式，進而利用這種複雜函式來擬合目標問題。

圖2. 增大網路容量常用的兩種方式

基於Transformer的研究，在機器翻譯任務上也取得了廣泛的關注，大致可分為兩類工作：第一類是透過結構最佳化的方式，讓網路中資訊的傳遞效率變得更高效；第二類主要針對深層網路中的引數初始化進一步改進。

眾所周知，目前 Transformer採用均勻初始化的方式得到w矩陣，該設計的核心思想是保證輸入輸出的方差一致性，最終得到了如圖3所示的邊界。然而伴隨網路深度的增加，對方差一致性的約束也變得更加薄弱，圖3中講者列舉的工作大部分針對如何在深層網路條件下保證方差一致性約束問題展開研究。

圖3. 在機器翻譯任務上深層Transformer的研究

綜上，深層Transformer模型具有顯著的效能優勢，但隨著網路深度增加（即模型越複雜），其計算量越大，佔用計算資源越多，難以直接應用於小裝置。那麼如何在保證效能的前提下，將複雜的深層網路壓縮至輕量級的網路？

知識蒸餾是解決上述問題的一種有效的技術手段，主要包括三種知識蒸餾方法：

第一種是基於詞級的方法，相當於在訓練student網路時，不僅能學習到ground truth，同時，透過增加teacher網路最頂層的表示與student表示之間的約束，以學習 teacher模型中的知識；

第二種區別於詞級方法，採用一種基於序列的方式，即在進行student網路的學習中，不學習ground truth而是學習給定原語的teacher模型輸出分佈。簡單來說將原語和teacher模型翻譯的譯文作為一個雙語句對，以供student模型學習；這種方式可以彌補student模型由於本身容量受限，對真實資料分佈加工能力不足的問題。

第三種是基於上述兩種方法的融合。

圖4. 常用知識蒸餾方法概述

圖5中的兩項工作致力於將BERT模型壓縮至一個輕量化的網路，其核心思想是選取 teacher模型和student模型中的attention分佈、詞嵌入表示、中間隱藏層狀態等，採用 MSE的方法構造損失函式。其弊端在於，在訓練過程中，要求同步計算student模型和teacher模型的分佈，消耗大量計算資源。因此，講者首先圍繞在機器翻譯任務上如何壓縮Transformer模型展開討論。

圖5. 關於壓縮BERT模型的相關工作

另一方面，在訓練student模型時，更好的teacher模型能夠提供更豐富的知識，通常採用增大teacher模型的深度來提高teacher模型的效能。但隨著網路深度的增加，模型更易陷入過擬合。

圖6. 模型深度增加導致過擬合現象

二、GPKD及子層跳躍方法介紹

針對第一個問題，如何更高效地學習輕量化的student模型？講者提出GPKD方法。相當於採用一種組內擾亂teacher的方法來訓練後，用teacher模型隱藏層狀態初始化student模型；然後透過序列化的知識蒸餾方法充分訓練student模型。

整個方法可以大致包括三個部分：

第一，進行組內擾亂的teacher訓練。首先，假設teacher有m層網路，student有n層網路，為將m層teacher的效能或知識遷移到 n層student模型上，用m除以n得到的數值，作為每個group內的層數；之後在訓練teacher模型時，每個group內的層計算順序可隨機打亂，且不同group間的順序擾亂是無關的；同時在不同batch的訓練過程中，不同batch間的擾亂也是獨立。該方式目的在於，讓每一個group內的每一個子層都能模仿整個group的輸出，便於後續student模型的訓練。

第二，採用標準的基於序列級的知識蒸餾方法。即在給定Src和Tgt雙語句對時，首先將得到的teacher模型進行 inference；然後將得到的SKD-Tgt以及Src作為 student模型的訓練資料。

第三，在訓練student的過程，隨機挑選每一個group內的一層網路引數來初始化student模型，該方法不僅能加快student模型的訓練，還能進一步逼近teacher效能，甚至達到和teacher相當效能。

圖7. GPKD方法介紹

針對第二個問題，如何解決模型過擬合？講者提出了子層跳躍的正則化訓練方法，即在進行網路前向計算時，採用dropout的思想對每一個子層進行控制，決定子層保留或丟棄。不同於常見dropout思路，在構造子層訓練方法時，採用自適應權重來區分不同深度子層被dropout的機率。該方式使得底層的子層更易被保留，而頂層的子層更易被丟棄，主要是由於，在層級連線網路中，底層網路的計算較頂層更加重要。在推斷階段，同樣基於dropout的思想，啟用所有分支，同時採用一種rescale方式，將每一個子層的輸出進行縮放。

圖8. 子層跳躍方法介紹

三、實驗

1.資料集介紹

實驗主要在WMT En-De，NIST Zh-En，WMT19 Zh-En三個任務上進行，其中訓練集、驗證集和測試集如表1所示。

表1. 實驗資料集介紹

2.訓練細則

訓練過程中採用Pre-Norm作為基線模型，分別對標準的Transformer模型以及加入相對位置編碼的RPR模型進行實驗。採用multi-bleu和sacrebleu分別作為NIST Zh-En和WMT19 Zh-En兩個任務上的評測指標。

3.GPKD在編碼端的實驗效果

圖a)藍色表示深層模型的效能，粉色表示6層基線模型的效能，由結果知：隨著網路深度增加，模型的效能顯著提升；綠色表示使用序列化知識蒸餾方法的效能，結果表明其能減小student模型與teacher模型之間的效能差距，但依舊存在較大的gap；黃色表示使用GPKD方法得到的student模型效能，結果表明在不同引數規模下， GPKD方法使student模型效能接近teacher模型，尤其從24層壓縮到6層時，幾乎沒有任何的效能損失。此外 GPKD方法較標準的序列化知識蒸餾方法有將近0.5-0.8幅度的效能提升。

圖b)結果表明，在NIST Zh-En資料集上，GPKD方法較SKD同樣有顯著的效能優勢。

a)WMT En-De

b) NIST Zh-En

圖9. GPKD在編碼端的效能對比實驗

4.GPKD在解碼端的實驗效果

進一步將該方法用於解碼端，分別衡量了隨著解碼端深度的增加，淺層和深層網路的BLUE效能以及推斷速度。圖中藍色的曲線表示深層網路，紅色曲線表示淺層的6層網路。由圖可知，隨著解碼端深度的增加，並沒有顯著的效能增益；反而在可接受效能範圍內，適度減小深層網路解碼層深度時，解碼端推斷速度可大幅度提高。講者將這類深層編碼器-淺層解碼器的組合稱之為異構網路，在一些工業應用上表現突出，既能保證翻譯的準確度，同時滿足工業生產的推斷時延需求。

圖10. GPKD在解碼端的效能對比實驗

5.子層跳躍方法的實驗效果

引入子層跳躍之後，模型在三個資料集上的校驗集PPL穩定下降，說明加入該方法後，有效提高了模型泛化能力，緩解過擬合現象。此外，在三個資料集上，網路效能也出現一定的提升，在WMT En-De資料集上，引入子層跳躍後較48層深度網路，BLUE效能提升了0.6。

圖11. 加入子層跳躍後的實驗驗證

圖12左圖表明，當在校驗集上呈現收斂趨勢後引入子層跳躍方法的網路，對比從頭訓練就使用子層跳躍方法的網路，具有顯著的效能優勢；同時與另外兩項工作相比，子層訓練方法較其在BLUE效能上存在較大競爭優勢。

針對如何選取子層跳躍的力度，進行了消融實驗。右圖結果表明，直接跳躍整層網路的計算，其效能略低於基線模型；隨機丟棄任何一個自注意機制網路或是前饋神經網路，其效能較基線模型有一定效能提升；同時丟棄自注意機制網路和前饋神經網路具有最好的效能。

圖12. 加入子層跳躍後的實驗驗證

四、總結

講者針對深層網路訓練中的兩個問題——如何更高效地學習輕量化的student模型以及如何緩解過擬合現象，分別提出了GPKD方法以及子層跳躍方法，結合兩種方法來增強teacher模型的訓練，進而得到一個更強的student模型。

相關資料

論文連結：

https://arxiv.org/pdf/2012.13866.pdf

Github地址：

https://github.com/libeineu/GPKD