沉重的 BERT

作者 | Derrick Mwiti

翻譯 | Camel

原文標題：Speeding up BERT

原文鏈接：

BERT 在 NLP 領域的地位正如 ResNet 在 CV 領域的地位一樣，屬於里程碑的進展。目前，BERT 已經成為 NLP 深度學習管道中的重要組成部分。

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

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

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

在基於 Pod 配置的 4 個 Cloud TPUs（總共 16 個 TPU 芯片）上對 BERT-base 進行訓練，或者在 16 個 Cloud TPU（總共

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

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

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

架構優化：將原有的架構改進為更快的架構，例如，將 RNN 替換為 Transformer 或 CNN；使用需要較少計算的層等。當然也可以採用其他優化，例如從學習率和策略、預熱步數，較大的批處理大小等；

模型壓縮：通常使用量化和修剪來完成，從而能夠在架構不變（或者大部分架構不變）的情況下減少計算總量；

模型蒸餾：訓練一個較小的模型，該模型能夠複製原始模型的行為。

我們來看下 BERT 在這些策略上可以做什麼。

1、架構和優化改進

大規模分佈式訓練

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

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

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

架 構

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

有一種基於對自注意層行為的觀察來訓練 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-large 具有相似配置的 ALBERT，其參數能夠降低 18 倍，訓練速度提高了月 1.7 倍。

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

2、量化和修剪

量化會降低模型權重的數值精度。通常情況下，使用 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」網絡的技術，訓練學生網絡來模仿教師網絡的行為。

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％。

比較 GLUE 基準測試的開發集

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

推理速度

TinyBERT

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

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

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

1、嵌入層的輸出；

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

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

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

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

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

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

TinyBERT 與其他基準的比較

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

其他蒸餾方法

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

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

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

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

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

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

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

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

TinyBERT 的結果似乎更好，但一個 7MB 的類 BERT 模型簡直爽的不能再爽！

需要強調，以上所介紹的方法並不互相沖突，所以我們期待能夠將這些方法協同起來使用來加速 BERT 模型（或其他相似模型）。

[1] https://www.alibabacloud.com/blog/perseus-bert-industry-leading-bert-training-solution-optimized-for-performance_594717

[2] https://github.com/NVIDIA/Megatron-LM

[3] https://devblogs.nvidia.com/training-bert-with-gpus/

[4] https://arxiv.org/abs/1904.00962

[5] http://proceedings.mlr.press/v97/gong19a/gong19a.pdf

[6] https://arxiv.org/abs/1904.10509

[7] https://arxiv.org/abs/1804.00857

[8] https://openreview.net/pdf?id=H1eA7AEtvS

[9] https://blog.inten.to/papers-roberta-a-robustly-optimized-bert-pretraining-approach-7449bc5423e7

[10] https://openreview.net/forum?id=HJe88xBKPr

[11] https://www.tensorflow.org/lite

[12] https://developer.nvidia.com/tensorrt

[13] https://news.developer.nvidia.com/tensorrt6-breaks-bert-record/

[14] https://pytorch.org/blog/pytorch-1-dot-3-adds-mobile-privacy-quantization-and-named-tensors/

[15]https://pytorch.org/tutorials/advanced/dynamic_quantization_tutorial.html

[16] https://arxiv.org/abs/1909.05840

[17] https://www.cs.cornell.edu/~caruana/compression.kdd06.pdf

[18] https://arxiv.org/abs/1503.02531

[19] https://arxiv.org/abs/1910.01108

[20] https://arxiv.org/abs/1909.10351

[21] https://arxiv.org/abs/1903.12136

[22] https://arxiv.org/abs/1908.09355

[23] https://github.com/intersun/PKD-for-BERT-Model-Compression

[24] https://arxiv.org/abs/1909.11687

閱讀更多 AI科技評論 的文章

關鍵字: 人工智能 算法 阿里巴巴集團