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

本文作者： camel

2019-10-28 18:56

導語：沉重的 BERT

原文作者：Grigory Sapunov

翻譯：Camel

原文標題：Speeding up BERT

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

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

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

沉重的BERT

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

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

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

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

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

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

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

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

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

1、架構和優(yōu)化改進

大規(guī)模分布式訓練

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

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

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

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

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

架構

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

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

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

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

ALBERT

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

ALBERT 結合了兩種參數(shù)降低的技術。

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

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

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

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

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

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

2、量化和修剪

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

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

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

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

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

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

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

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

更多的數(shù)據(jù)我們這里就不再一一列舉了。

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

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

3、蒸餾

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

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

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

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

DistilBERT

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

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

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

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

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

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

比較 GLUE 基準測試的開發(fā)集

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

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

推理速度

TinyBERT

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

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

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

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

1、嵌入層的輸出；

2、Transformer 層派生的隱藏狀態(tài)和注意力矩陣;

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

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

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

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

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

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

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

TinyBERT 與其他基準的比較

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

其他蒸餾方法

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

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

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

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

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

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

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

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

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

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

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