Feature map transfer in TinyBert/MobileBert


MobileBert: 自下而上的知識轉移: 將 Bottleneck BERT-LARGE teacher 的知識傳遞給student。

先訓練teacher,然后逐步從下網上訓練student,要求它一層一層地模仿teacher。

 

知識轉移的漸進過程分為L個階段,其中L為層數。圖2展示了漸進式知識轉移的原理圖和算法。漸進遷移的思想是在訓練學生的(l+1)層時,(l)層已經是最優的了。

由於BERT的中間狀態沒有軟目標分布,我們提出了兩個知識轉移目標: feature map transfer和attention transfer來訓練學生網絡。特別地,我們假設老師和學生有相同的1)feature map的大小,2)層的數量,3)attention heads的數量。

 

FEATURE MAP TRANSFER (FMT)

由於BERT中的每一層僅僅是將前一層的輸出作為輸入,所以在逐步訓練學生網絡時,最重要的是每一層的feature maps應該盡可能接近於教師的feature maps,即優化的。其中,以學生與教師歸一化特征圖的均方誤差為目標:

其中l為層索引,T為序列長度,N為特征映射大小。增加層間歸一化以穩定層間訓練損失。
我們還最小化了特征圖轉換中均值和方差的兩個統計差異:

 

µ和σ2代表均值和方差。實驗表明,在BERT中去除層標准化以減少推理延遲時,最小化統計差異是有幫助的(參見4.3節中的更多討論)。 

 

ATTENTION TRANSFER (AT)

BERT中的注意分布可以檢測出詞語之間合理的語義和句法關系。

這激發了我們使用來自優化過的教師的“自我注意maps”來幫助訓練學生在增強到“feature map transfer”的過程中。特別是,我們最小化了教師和學生的pre-head的(head平均的)注意力分布的kl-散度:

 

 A是attention heads的數目。

 

 

TinyBert

作者提出了兩個點:針對Transformer結構的知識蒸餾和針對pre-training和fine-tuning兩階段的知識蒸餾。

作者在這里構造了四類損失函數來對模型中各層的參數進行約束來訓練模型,具體模型結構如下:

    

   作者構造了四類損失,分別針對embedding layer,attention 權重矩陣,隱層輸出,predict layer。可以將這個統一到一個損失函數中:

    

   上面式子中𝜆𝑚λm表示每一層對應的系數,𝑆𝑚Sm表示studnet網絡的第m層,𝑇𝑔(𝑚)Tg(m)表示teacher網絡的第n層,其中𝑛=𝑔(𝑚)n=g(m)。並且有𝑔(0)=0g(0)=0,𝑔(𝑀+1)=𝑁+1g(M+1)=N+1,0表示embedding layer,M+1和N+1表示perdict layer。

  針對上面四層具體的損失函數表達式如下:

  attention 權重矩陣

    

   h為multi attention中頭數

  隱層輸出

    

   因為student網絡的隱層大小通常會設置的比teacher的小,因此為了在計算時維度一致,這里用一個矩陣𝑊Wh將student的隱層向量線性映射到和teacher同樣的空間下。

  embedding layer

    

   𝑊𝑠Ws同理上。

  以上三種損失函數都采用了MSE,主要是為了將模型的各項參數對齊。

  predict layer

    

   predict layer也就是softmax層,在這里的損失函數是交叉熵,t是溫度參數,在這里設置為1。

  以上四種損失函數是作者針對transformer提出的知識蒸餾方法。除此之外作者認為除了對pre-training蒸餾之外,在fine-tuning時也利用teacher的知識來訓練模型可以取得在下游任務更好的效果。因此作者提出了兩階段知識蒸餾,如下圖所示:

    

 

   本質上就是在pre-training蒸餾一個general TinyBERT,然后再在general TinyBERT的基礎上利用task-bert上再蒸餾出fine-tuned TinyBERT。

  作者給出了TinyBERT的效果:

    

   另外作者也給出了四種損失對最終結果的貢獻:

    

   還有就是關於𝑛=𝑔(𝑚)n=g(m)這個式子中𝑔(𝑚)g(m)怎么選擇,假設student的層數為4層,這里的𝑛=𝑔(𝑚)=3𝑚n=g(m)=3m,作者將這種稱為Uniform-strategy。另外作者還和其他的𝑔(𝑚)g(m)做了對比:

    

   Top-strategy指用teacher最后4層,Bottom-strategy指用前面4層,其實這里的映射函數,我感覺可能還有更優的方案,例如取平均,或者用attention來做,可能效果會更好。

 

 

 


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM