加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

本文作者： camel

2019-10-28 18:56

導語：沉重的 BERT

原文作者：Grigory Sapunov

翻譯：Camel

原文標題：Speeding up BERT

原文鏈接：https://blog.inten.to/speeding-up-bert-5528e18bb4ea

雷鋒網AI科技評論按：BERT 在 NLP 領域的地位正如 ResNet 在 CV 領域的地位一樣，屬于里程碑的進展。目前，BERT 已經成為 NLP 深度學習管道中的重要組成部分。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

沉重的BERT

但 BERT 并不是完美無瑕的，它最大的問題就是：太大。

BERT-base 模型能夠包含一億個參數，較大的 BERT-large 甚至包含 3.4 億個參數。顯然，很難將這種規模的模型部署到資源有限的環境（例如移動設備或嵌入式系統）當中。

模型太大是其一，BERT 的訓練和推理時間也太長了！

在基于 Pod 配置的 4 個 Cloud TPUs（總共 16 個 TPU 芯片）上對 BERT-base 進行訓練，或者在 16 個 Cloud TPU（總共 64 個 TPU 芯片）上對 BERT-large 進行訓練，每次預訓練都需要至少 4 天的時間才能完成。

當然對于最終的用戶而言，訓練時間通常似乎并不是什么大問題，因為反正只需要訓練一次就夠了（實際上往往會訓練多次，因為你可能要多次重新訓練模型，才能獲得滿意的結果）。不過如果能夠加速，當然也是好的，因為迭代越快，你就能越早地解決你的問題。

BERT 的推理時間可能取決于你用的是什么設備（什么型號、多少個），在許多情況下，它會極大地限制你要處理數據的數量、速度以及成本。對于某些實時的應用程序來講，用 BERT 簡直是讓人望而卻步。

對以上的問題，能夠想到的最直接的方法就是優化神經網絡。這在神經網絡中并不陌生，其他領域（例如計算機視覺）以前也有相同的問題，并且已經開發了幾種壓縮和加速神經網絡模型的方法。大致可以分為幾類：

架構改進：將原有的架構改進為更快的架構，例如，將 RNN 替換為 Transformer 或 CNN；使用需要較少計算的層等。當然也可以采用其他優化，例如從學習率和策略、預熱步數，較大的批處理大小等；
模型壓縮：通常使用量化和修剪來完成，從而能夠在架構不變（或者大部分架構不變）的情況下減少計算總量；
模型蒸餾：訓練一個較小的模型，該模型能夠復制原始模型的行為。

我們來看下 BERT 在這些策略上可以做什么。雷鋒網

1、架構和優化改進

大規模分布式訓練

加速 BERT 訓練能夠想到的第一件事（甚至是第零件事）是將其分發到更大的集群當中。雖然最初的 BERT 已經使用了多臺機器進行訓練，但還有更多針對 BERT 分布式訓練的優化解決方案，例如阿里巴巴 [1] 和英偉達 [2] 的一些工作。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

英偉達最近使用 NVIDIA DGX SuperPOD（具有 92 個 DGX-2H 節點，共有 1472 個 V100 GPU，理論上可以提供 190PFLOPS）刷新了 BERT 訓練的記錄，在 53 分鐘內訓練出了 BERT-large 模型 [3]。當然，這個訓練也是非常昂貴的，除了英偉達，其他公司和個人很難去做這樣的事情。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

還有另外一種更為巧妙優化方案，這是一種被稱為 LAMB 的新的分層自適應大批量優化技術 [4]，這種技術可以將 TPUv3 Pod 上的 BERT 訓練時間從 3 天降低到 76 分鐘（1024 個 TPUv3 芯片，可以提供超過 100PFLOPS），顯然，訓練這個模型也不會便宜。

架構

現在考慮更加架構性的以及更少硬件的解決方案。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

有一種基于對自注意層行為的觀察來訓練 BERT 的漸進式堆疊方法 [5]，該方法的分布局部地集中在其位置和 start-of-sentence token 周圍，且在淺層模型中的注意力分布與深層模型相似。為此，作者提出了將知識從淺層模型轉移到深層模型的堆疊算法，并逐步應用堆棧來加快 BERT 訓練的速度。憑借這種方式，作者獲得模型的訓練時間相比原始 BERT 的訓練時間縮短了約 25%，歸其原因則在于對于相同數量的步驟，訓練一個小的模型需要的計算也相對較少。

還有一種方法是注意力矩陣的稀疏分解（也稱 Sparse Transformer，這項研究是由 OpenAI 完成的)[6] 和 block attention[7]，它們都屬于能夠減少內存和計算總量的體系結構改進。

ALBERT

最后不得不提的是，已經提交給 ICLR 2020 的一個 BERT 的后代，被稱為 ALBERT（A Lite BERT）[8]。

ALBERT 結合了兩種參數降低的技術。

第一種是分解式嵌入參數化，將隱藏層的大小與詞嵌入的大小分開。這種分隔使得在不顯著增加詞匯表嵌入參數大小的情況下能夠更容易增加隱藏層的大小。

第二種是跨層參數共享。這種技術可以防止參數隨著網絡深度的增加而增大。

這兩種技術可以顯著降低 BERT 的參數數量，而不會嚴重影響性能，從而提高參數效率。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

與 BERT-large 具有相似配置的 ALBERT，其參數能夠降低 18 倍，訓練速度提高了月 1.7 倍。

相比嚴格調校的 RoBERTa，ALBERT 則同樣更勝一籌 [9]。

2、量化和修剪

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

量化會降低模型權重的數值精度。通常情況下，使用 FP32（32 位浮點）來訓練模型，然后可以將其量化為 FP16（16 位浮點），INT8（8 位整數），甚至將其量化為 INT4 或 INT1。于是模型尺寸將隨之減小 2 倍、4 倍、8 倍或 32 倍。這稱為 post-training quantization。

另一個選項是量化感知訓練（也是較為困難和較為不成熟的）。這種方法的 FP16 訓練現在已成為一種普遍的方法。而在 ICLR 2020 提交的論文中有一個比較有趣的工作 [10]，它使用 ResNet、GNMT、Transformer 的 8 位浮點表示獲得了最先進的訓練結果。

修剪即從模型中刪除一些不重要或不太重要的權重（有時會是神經元），從而產生稀疏的權重矩陣（或較小的圖層）。甚至還有一些研究直接去除掉與 Transformer 的注意力頭相對應的整個矩陣。

量化，可以使用 TensorFlow Lite（用于設備上推斷的 TensorFlow 的部分）[11] 來執行。TensorFlow Lite 提供了在移動設備、嵌入式設備和 IoT 設備上轉換和運行 TensorFlow 模型的工具，它支持訓練后量化和量化感知訓練。

另一個選擇是使用英偉達的 TensorRT 框架 [12]。英偉達 TensorRT 是用于高性能深度學習推理的平臺，它包括深度學習推理優化器，并且在運行時能夠為深度學習推理應用程序提供低延遲和高吞吐量。

英偉達最近發布了具有最新優化功能的 TensorRT 6[13]，它可以在 T4 GPU 上僅用 5.8 毫秒對 BERT-large 進行推理，而在 V100 上甚至只需要 4.2 毫秒，對于 Titan RTX，速度可能會更快。

當批處理數為 1 時，對于 340M 個參數的 BERT-large 模型，僅需 5.84 毫秒；對于 110M 參數的 BERT-Base 模型，則只需 2.07 毫秒。若將批處理數設為 128，你在使用 BERT-large 時甚至可以達到 250 個句子/秒的處理速度。

更多的數據我們這里就不再一一列舉了。

PyTorch 最近也宣布了在其 1.3 版本 [14] 中支持量化。盡管目前它還處于實驗階段，但我們已經可以使用它了，因為在其教程中提到他們已經能夠將動態量化應用于將模型權重轉換為 INT8 的 LSTM 語言模型 [15]。

此外，還有一個眾所周知的 BERT 量化，稱為 Q-BERT。這項工作來源于「Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT」[16]。在這個研究中，作者甚至降低到 2-bit 的超低精度量化，但其性能相比于基線卻沒有顯著下降（僅下降 2.3%），而對應的模型參數壓縮率最高可以達 13 倍，嵌入表壓縮率和激活的最高都為 4 倍。

3、蒸餾

另一個有趣的模型壓縮方法是蒸餾，這是一種將大型「teacher」網絡的知識轉移到較小的「student」網絡的技術，訓練學生網絡來模仿教師網絡的行為。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

Rich Caruana 及其合作者率先采用了這種策略。在他們先驅性的論文 [17] 中，他們提供了令人信服的證明：大型集成模型所獲得的知識可以轉移到單個小型的模型中。

Geoffrey Hinton 等人在他們的「Distilling the Knowledge in a Neural Network」{18} 論文中證明了這種技術可以應用于神經網絡模型。

DistilBERT

從 Hinton 開始，蒸餾的方法逐漸被應用到了不同的神經網絡中，當然你可能也聽過 HuggingFace 提出的 DistilBERT，這是一種對 BERT 的蒸餾。這項工作出自論文「DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter」[19]，目前已經提交到 NeurIPS 2019。

DistilBERT 是一種較小的語言模型，受 BERT 的監督而訓練。在該模型中，作者刪除了令牌類型嵌入和合并器（用于下一個句子分類任務），并保持體系架構其余部分不變，同時將層數減少了兩倍。

您可以在 HuggingFace（以前叫做 pytorch-transformers 和 pytorch-pretrained-bert）的 translators python 軟件包的幫助下，使用現成的 DistilBERT。該軟件包的 2.0.0 版本支持 TensorFlow 2.0 / PyTorch 互操作。

DistilBERT 的作者還使用了最新 RoBERTa 論文中的一些訓練技巧，這些技巧表明 BERT 的訓練方式對其最終性能至關重要。

DistilBERT 與 BERT 相比具有出人意料的結果：作者保留了 95％以上的性能，但參數卻減少了 40％。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

比較 GLUE 基準測試的開發集

在推斷時間方面，DistilBERT 比 BERT 快了 60％以上，比 ELMo + BiLSTM 快 120％。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

推理速度

TinyBERT

就在幾天前，出現了一種新的 BERT 蒸餾方法，來自華為諾亞方舟實驗室的劉群團隊提出了 TinyBERT[20]。

為了構建一個具有競爭力的 TinyBERT，作者首先提出了一種新的 Transformer 蒸餾方法，來蒸餾嵌入 BERT 的知識。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

具體來說就是，他們設計了幾個損失函數來適合 BERT 層的不同表示形式：

1、嵌入層的輸出；

2、Transformer 層派生的隱藏狀態和注意力矩陣;

3、預測層輸出的 logits 輸出。

論文中基于注意力的擬合則得益于最近的發現，即 BERT 學習的注意力權重可以捕獲大量的語言知識，這意味著語言知識可以很好地從教師 BERT 轉移到學生 TinyBERT。而在 BERT 的現有知識蒸餾的方法（如 Distilled BiLSTM_SOFT，BERT-PKD 和 DistilBERT）中卻忽略了這一點。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

在這項工作中，作者還提出了一種新穎的兩階段學習框架，包括通用蒸餾和特定任務蒸餾。在通用蒸餾階段，未經微調的原始 BERT 充當教師模型，學生 TinyBERT 通過在通用領域對大型語料庫執行通常的 Transformer 蒸餾來學習模仿教師的行為。他們獲得了可以對各種下游任務進行微調的通用 TinyBERT。在特定任務蒸餾階段，他們將數據進行擴充，來提供更多與任務相關的材料供教師-學生學習，然后在增強的數據上重新執行 Tranformer 蒸餾。

這個兩階段方法對于提升 TinyBERT 的性能和泛化能力是必不可少的。

TinyBERY 在實驗上取得了非常的成績，相對于 GLUE 數據集的 BERT-base，其性能并沒有下降多少，而推理參數小了 7.5 倍，推理時間快了 9.4 倍。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

TinyBERT 與其他基準的比較

我們期待他們能夠將這種方法應用到 BERT-large 和 XLNet 等大型模型中，同樣也期待他們開放源碼。

其他蒸餾方法

除了 DistilBERT 和 TinyBERT 外，還有其他一些為大家所熟知的蒸餾方法。

（2019/03）「Distilling Task-Specific Knowledge from BERT into Simple Neural Networks」[21]

這篇論文將 BERT 蒸餾到單層 BiLSTM 中，取得了與 ELMo 可比的結果，同時使用的參數減少了大約 100 倍，推理時間減少了 15 倍。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

BiLSTM_SOF 是 TBiLSTM 的蒸餾，后者是在 soft logit 目標上訓練出來的。

（2019/08）「Patient Knowledge Distillation for BERT Model Compression」[22]

這篇論文提出了一種耐心知識蒸餾的方法，這是首次嘗試使用教師的隱藏狀態，而不僅僅是最后一層的輸出。他們的學生模型從教師模型的多個中間層「耐心」地學習來獲得更多知識。在他們的耐心蒸餾知識框架中，只訓練學生模仿中間層的 [CLS] 令牌的表示形式。代碼已公開 [23]。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！

（2019/09）「Extreme Language Model Compression with Optimal Subwords and Shared Projections」

這是最近提交到 ICLR 2020 的一篇論文，這篇論文專注于一種用于訓練詞匯量顯著較小、嵌入和隱藏狀態維度較低的學生模型的知識蒸餾技術。作者采用了雙重訓練機制，可以同時訓練教師和學生模型，從而能夠獲得針對學生詞匯的最佳詞嵌入。該方法能夠將 BERT-base 模型壓縮 60 倍以上，而下游任務指標只有很小的下降，從而使得語言模型占用的空間只有不到 7MB。

加速BERT模型有多少種方法？從架構優化、模型壓縮到模型蒸餾，最新進展詳解！