哈工大訊飛聯合實驗室（HFL）在前期釋出了多箇中文預訓練系列模型，目前已成為最受歡迎的中文預訓練資源之一。然而，眾多預訓練模型體積龐大，難以滿足執行時要求，為技術落地提出了新的挑戰。為此，我們很高興地向大家推出基於PyTorch框架的知識蒸餾工具包TextBrewer，提供更加方便快捷的知識蒸餾框架。我們歡迎各位讀者積極試用並反饋寶貴意見。
工具地址：http://textbrewer.hfl-rc.com
論文地址：https://arxiv.org/abs/2002.12620
特點
TextBrewer為NLP中的知識蒸餾任務設計，提供方便快捷的知識蒸餾框架，主要特點包括：
- 模型無關：適用於多種模型結構（主要面向Transfomer結構）
- 方便靈活：可自由組合多種蒸餾方法，支援增加自定義損失等模組
- 非侵入式：無需對教師與學生模型本身結構進行修改
- 適用面廣：支援典型NLP任務，如文字分類、閱讀理解、序列標註等
TextBrewer目前支援的主要知識蒸餾技術有：
- 軟標籤與硬標籤混合訓練
- 動態損失權重調整與蒸餾溫度調整
- 多種蒸餾損失函式
- 任意構建中間層特徵匹配方案
- 多教師知識蒸餾
- ...
工作流程
TextBrewer工具中的一個完整工作流程如下圖所示。
TextBrewer整體工作流程
△ 第一步：在開始蒸餾之前的準備工作
- 訓練教師模型
- 定義並初始化學生模型（隨機初始化或載入預訓練權重）
- 構造蒸餾用資料集的DataLoader，訓練學生模型用的Optimizer和Learning rate scheduler
△ 第二步 : 知識蒸餾
- 初始化Distiller，構造訓練配置（TrainingConfig）和蒸餾配置（DistillationConfig）
- 定義adaptors和callback，分別用於適配模型輸入輸出和訓練過程中的回撥
- 呼叫Distiller的train方法開始蒸餾
使用者應先進行第一步準備工作，得到訓練好的教師模型。TextBrewer負責第二步的知識蒸餾工作。為了方便使用者使用，TextBrewer也提供了BasicTrainer用於訓練第一步的教師模型。
知識蒸餾本質上是“老師教學生”的過程。在初始化學生模型時，可以採用隨機初始化的形式（即完全不包含任何先驗知識），也可以載入已訓練好的模型權重。例如，從BERT-base模型蒸餾到3層BERT時，可以預先載入RBT3模型權重，然後進一步進行蒸餾，避免了蒸餾過程的“冷啟動”問題。我們建議使用者在使用時儘量採用已預訓練過的學生模型，以充分利用大規模資料預訓練所帶來的優勢。
實驗效果
為了驗證工具的效果，我們在多箇中英文自然語言處理任務上進行了實驗並取得了與業界公開結果相當，甚至是超過相關公開工作的效果。我們對如下幾種基於Transformer的模型結構進行了知識蒸餾。
知識蒸餾實驗中使用的模型結構和引數量在中文實驗中，我們選取了4個經典資料集，其中包括XNLI（自然語言推斷）、LCQMC（句對分類）、CMRC 2018（閱讀理解）、DRCD（繁體閱讀理解）。其中教師模型為哈工大訊飛聯合實驗室釋出的RoBERTa-wwm-ext模型。
中文任務上的知識蒸餾效果在英文實驗中，我們選取了3個經典資料集，其中包括MNLI（自然語言推斷）、SQuAD（閱讀理解）、CoNLL-2003 NER（命名實體識別）。我們也列出了一些公開論文中的知識蒸餾效果，方便大家進行比較。教師模型為谷歌釋出的BERT-base-cased模型。
6層Transformer級別上的效果對比
3層Transformer級別上的效果對比
與TinyBERT的效果對比
相關資源地址
- TextBrewer知識蒸餾工具
- http://textbrewer.hfl-rc.com
- 中文BERT、RoBERTa、RBT系列模型
- http://bert.hfl-rc.com
- 中文XLNet系列模型
- http://xlnet.hfl-rc.com