小版BERT也能出奇跡：最火的預訓練語言庫探索小巧之路

選自Medium

作者：Victor Sanh

機器之心編譯

參與：魔王

過去一段時間，大模型層出不窮。在大家紛紛感歎「大力出奇跡」的時候，作為調用預訓練語言模型最流行的庫，HuggingFace 嘗試用更少的參數量、更少的訓練資源實現同等的性能，於是知識蒸餾版 BERT——DistilBERT 應運而生！

近日，HuggingFace 發布了 NLP transformer 模型——DistilBERT，該模型與 BERT 的架構類似，不過它僅使用了 6600 萬參數（區別於 BERT_base 的 1.1 億參數），卻在 GLUE 基準上實現了後者 95% 的性能。

在過去 18 個月中，基於大規模語言模型的遷移學習顯著提升了自然語言處理任務的 SOTA 成績。這些預訓練語言模型通常基於 Vaswani 等人提出的 Transformer 架構，這些模型的規模越來越大，訓練數據集的規模也越來越大。最近NVIDIA提出的預訓練語言模型擁有 83 億參數：是 BERT-large 參數量的 24 倍、GPT-2 參數量的 5 倍。而 Facebook AI 最近提出的 RoBERTa 模型在 160GB 文本上訓練得到。

社區中的一些人質疑訓練越來越大 Transformer 的必要性，尤其是考慮到訓練的資金成本和環境成本時。該圖展示了部分近期大模型及其參數量。

Hugging Face 直接體會到這些模型的流行度，因為其預訓練語言庫（包含這些模型中的大部分）在近幾個月的安裝量超過 40 萬次。

NLP 庫地址：https://github.com/huggingface/pytorch-transformers

然而，儘管這些模型被更大的 NLP 社區接受，一個重要且有挑戰性的問題出現了。如何將這些龐然大物投入到生產中？如何在低延遲約束下使用這些大模型？我們需要用（昂貴的）GPU 伺服器執行大規模服務嗎？

為了構建更尊重隱私的系統，Hugging Face 注意到在邊緣設備上運行機器學習系統的需求在不斷增長，而不是調用雲 API，將隱私數據發往伺服器。在智能手機等設備上運行的模型需要是輕量級、響應快和能源利用率高的！

最後但同樣重要的一點，Hugging Face 越來越擔憂這些大模型所需的指數級計算成本增長。

有很多技術可以解決前述問題。最常見的工具是量化（使用更小精度逼近全精度模型）和權重剪枝（移除網絡中的部分連接）。想了解更多，可以查看這篇關於 BERT 量化的精彩部落格：https://blog.rasa.com/compressing-bert-for-faster-prediction-2/。

Hugging Face 的研究者決定把重點放在知識蒸餾（distillation）上。蒸餾即，將較大模型（教師模型）壓縮成較小模型（學生模型）的方法。

知識蒸餾：遷移泛化能力

知識蒸餾是一種模型壓縮方法，又叫師生學習（teacher-student learning）。它訓練一個小模型，使之複製大模型（或模型集成）的行為。知識蒸餾由 Bucila 等人提出，幾年後被 Hinton 等人推廣（參見論文《Distilling the Knowledge in a Neural Network》）。Hugging Face 研究者使用的是 Hinton 等人的方法。

在監督學習中，分類模型通常用於預測類別，它利用對數似然信號最大化類別概率。在很多案例中，高性能模型預測的輸出分布中，正確的類別具備高概率，而其他類別的概率則接近於零。

例如，desk chair（辦公椅）可能會被誤分類為 armchair（扶手椅），但通常不會被誤認為是 mushroom（蘑菇）。這種不確定性被稱為「暗知識」。

理解蒸餾的另一種方式是，它阻止模型對預測結果過於自信（類似於標簽平滑）。

以下是一個實踐示例。在語言建模過程中，我們通過觀察詞匯分布，可以輕鬆發現這種不確定性。下圖展示了 Bert 對電影《卡薩布蘭卡》中某句著名台詞的 top 20 補全結果：

BERT_base 對被遮蔽 token 的 top 20 補全結果。該語言模型確定了兩個概率較高的 token（day 和 life）。

如何複製暗知識？

在師生訓練中，我們訓練學生網絡來模擬教師網絡的完整輸出分布（它的知識）。

我們使學生網絡和教師網絡具備同樣的輸出分布，從而使學生網絡實現同樣的泛化。

我們不對硬目標類別（正確類別的 one-hot 編碼）使用交叉熵來進行訓練，而是對軟目標類別（教師網絡的概率）執行交叉熵，從而將教師網絡的知識遷移到學生網絡。這樣訓練損失函數就變成了：

其中 t 表示教師網絡的 logit 值，s 表示學生網絡的 logit 值。該損失函數具備更豐富的訓練信號，因為軟目標類別比單個硬目標類別提供更多約束。

為了進一步揭示類別分布的多樣性，Hinton 等人提出了 softmax-temperature：

其中T 表示溫度參數，當 T 0 時，分布接近於 one-hot 目標向量，當 T +∞ 時，則得到均勻分布。在訓練過程中對教師網絡和學生網絡使用同樣的溫度參數，進而為每一個訓練樣本提供更多信號。在推斷時，T 被設置為 1，恢復標準的 Softmax 函數。

PyTorch 動手實踐：壓縮 BERT

Hugging Face 研究者想利用知識蒸餾壓縮大型語言模型。對於蒸餾，研究者使用 KL 散度作為損失函數，因為最優化過程與交叉熵是等價的：

在計算 q（學生網絡的分布）的梯度時，得到了同樣的梯度。這允許研究者利用 PyTorch 實現執行更快速的計算：

使用教師網絡 BERT 的監督信號，研究者訓練得到較小的語言模型——DistilBERT。（研究者使用的是 Bert 的英語 bert-base-uncased 版本）。

按照 Hinton 等人的方法，訓練損失是蒸餾損失和遮蔽語言建模損失的線性組合。學生模型是 BERT 的較小版本，研究者移除了 token 類型的嵌入和 pooler（用於下一句分類任務），保留了 BERT 的其余架構，不過網絡層數只有原版的 1/2。

備注 1：為什麽不減少隱藏層大小呢？將它從 768 減少到 512 即可將參數總量減少約 1/2。但是，在現代框架中，大部分運算是經過高度優化的，張量最後一維（隱藏維度）的變化對 Transformer 架構中使用的大部分運算影響較小。在研究者的實驗中，相比隱藏層大小，層數才是推斷階段的決定性因素。

研究者的早期實驗表明，在該案例中，交叉熵損失會帶來更好的性能。因此，他們假設在語言建模設置中，輸出空間（詞匯）要比下遊任務輸出空間的維度大得多。而在 L2 損失中，logit 可能會相互抵消。

訓練子網絡不僅僅關乎架構，它還需要找出子網絡收斂的合適初始化（例如彩票假設論文《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks》）。因此，研究者基於教師網絡 Bert 對學生網絡 DistilBERT 進行初始化，將層數減半，對學生網絡使用與教師網絡一樣的隱藏層大小。

研究者還使用 RoBERTa 論文中的一些訓練 trick，這篇論文證明 Bert 的訓練方式對最終性能有著重要影響。遵循 RoBERTa 的訓練方式，研究者利用梯度累積以非常大的批次（每個批次多達 4000 個樣本）訓練 DistilBERT，訓練使用了動態遮擋（dynamic masking），並移除了下一句預測目標。

該訓練設置主動接受資源方面的限制。研究者在 8 塊 16GB V100 GPU 上訓練 DistilBERT，訓練時長接近三天半，訓練數據為 Toronto Book Corpus 和英文維基百科（與原版 BERT 的訓練數據相同）。

DistilBert 的部分代碼來自於 Facebook XLM 的代碼，部分代碼來自 Google AI BERT 的 Hugging Face PyTorch 版本。這些代碼可在 Hugging Face 的 NLP 庫中獲取，該庫還包含多個 DistilBert 訓練版本和微調版本，及其複現代碼。

模型性能：DistilBERT 測試

研究者在 GLUE 基準的開發集上對比了 DistilBERT 和兩個基線模型的性能，基線模型分別是 BERT base（DistilBERT 的教師模型）和來自紐約大學的強大非 transformer 基線模型：ELMo + BiLSTMs。研究者使用紐約大學發布的 ELMo 基線 jiant 庫和 BERT 基線模型的 PyTorch-Transformers 版本。

如下表所示，DistilBERT 在參數量分別是基線模型的 1/2 和 1/3 的情況下，性能可與後者媲美。在 9 項任務中，DistilBERT 的性能通常等同於或優於 ELMo 基線（在 QNLI 任務上的準確率超出後者 14 個百分點）。令人驚訝的是，DistilBERT 的性能堪比 BERT：在參數量比 BERT 少 40% 的情況下，準確率達到了後者的 95%。

在 GLUE 基準開發集上的模型對比結果。ELMo 的性能結果來自原論文，BERT 和 DistilBERT 的性能結果是使用不同種子進行 5 次運行後的中位數。

至於推斷性能，DistilBERT 的推斷速度比 BERT 快 60%，規模也比後者小；DistilBERT 的推斷速度比 ELMo+BiLSTM 快 120%，規模也比後者小很多。

為了進一步調查 DistilBERT 的加速/規模權衡（speed-up/size trade-off），研究者對比了每個模型的參數量和在 STS-B 開發集上使用一塊 CPU、批大小為 1 的情況下完成一個完整 pass 的推斷時間，如上表所示。

下遊任務：蒸餾和遷移學習

研究者進一步研究了在高效推斷約束下，DistilBERT 在下遊任務上的應用。研究者對緊湊的預訓練語言模型 DistilBERT 進行微調，用於分類任務。這是結合蒸餾預訓練和遷移學習的絕妙方式！

研究者使用 IMDB 評論情感分類數據集，該數據集包含 5 萬條英文評論（被標注為積極或消極）：其中 2.5 萬條作為訓練數據，另外一半作為測試數據（均類別均衡）。研究者使用一塊 12GB K80 GPU 進行訓練。

首先，在數據集上訓練 bert-base-uncased。該模型達到了 99.98% 的準確率（3 次運行的平均值），幾乎完美！

然後使用同樣的超參數訓練 DistilBERT。該模型達到了 99.53% 的準確率（3 次運行的平均值），在延遲降低 60%、規模減少 40% 的情況下，DistilBERT 的性能僅比原版 BERT 低 0.5%！

少即是多：小模型也能出奇跡

Hugging Face 對 DistilBERT 的潛力非常看好。DistilBERT 只是個開始，它也提出了許多問題：使用知識蒸餾技術，我們可以把大模型壓縮到什麽程度？這些技術可用於進一步探索和洞察大模型中存儲的知識嗎？在壓縮過程中損失了語言學/語義學的哪些方面？……

目前，HuggingFace 的這項研究已經開源，並進行知識共享。他們認為這是每個人參與 NLP 進展，並收獲最新進展果實的最快和最公平的路徑。

GitHub 地址：https://github.com/huggingface

Medium 頁面：http://www.medium.com/huggingface

本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。

------------------------------------------------