原文作者：Grigory Sapunov
翻譯：Camel
原文標(biāo)題：Speeding up BERT
原文鏈接：https://blog.inten.to/speeding-up-bert-5528e18bb4ea
雷鋒網(wǎng)AI科技評(píng)論按：BERT 在 NLP 領(lǐng)域的地位正如 ResNet 在 CV 領(lǐng)域的地位一樣，屬于里程碑的進(jìn)展。目前，BERT 已經(jīng)成為 NLP 深度學(xué)習(xí)管道中的重要組成部分。
沉重的BERT
但 BERT 并不是完美無瑕的，它最大的問題就是：太大。
BERT-base 模型能夠包含一億個(gè)參數(shù)，較大的 BERT-large 甚至包含 3.4 億個(gè)參數(shù)。顯然，很難將這種規(guī)模的模型部署到資源有限的環(huán)境（例如移動(dòng)設(shè)備或嵌入式系統(tǒng)）當(dāng)中。
模型太大是其一，BERT 的訓(xùn)練和推理時(shí)間也太長了！
在基于 Pod 配置的 4 個(gè) Cloud TPUs（總共 16 個(gè) TPU 芯片）上對(duì) BERT-base 進(jìn)行訓(xùn)練，或者在 16 個(gè) Cloud TPU（總共 64 個(gè) TPU 芯片）上對(duì) BERT-large 進(jìn)行訓(xùn)練，每次預(yù)訓(xùn)練都需要至少 4 天的時(shí)間才能完成。
當(dāng)然對(duì)于最終的用戶而言，訓(xùn)練時(shí)間通常似乎并不是什么大問題，因?yàn)榉凑恍枰?xùn)練一次就夠了（實(shí)際上往往會(huì)訓(xùn)練多次，因?yàn)槟憧赡芤啻沃匦掠?xùn)練模型，才能獲得滿意的結(jié)果）。不過如果能夠加速，當(dāng)然也是好的，因?yàn)榈娇?，你就能越早地解決你的問題。
BERT 的推理時(shí)間可能取決于你用的是什么設(shè)備（什么型號(hào)、多少個(gè)），在許多情況下，它會(huì)極大地限制你要處理數(shù)據(jù)的數(shù)量、速度以及成本。對(duì)于某些實(shí)時(shí)的應(yīng)用程序來講，用 BERT 簡(jiǎn)直是讓人望而卻步。
對(duì)以上的問題，能夠想到的最直接的方法就是優(yōu)化神經(jīng)網(wǎng)絡(luò)。這在神經(jīng)網(wǎng)絡(luò)中并不陌生，其他領(lǐng)域（例如計(jì)算機(jī)視覺）以前也有相同的問題，并且已經(jīng)開發(fā)了幾種壓縮和加速神經(jīng)網(wǎng)絡(luò)模型的方法。大致可以分為幾類：
架構(gòu)改進(jìn)：將原有的架構(gòu)改進(jìn)為更快的架構(gòu)，例如，將 RNN 替換為 Transformer 或 CNN；使用需要較少計(jì)算的層等。當(dāng)然也可以采用其他優(yōu)化，例如從學(xué)習(xí)率和策略、預(yù)熱步數(shù)，較大的批處理大小等；
模型壓縮：通常使用量化和修剪來完成，從而能夠在架構(gòu)不變（或者大部分架構(gòu)不變）的情況下減少計(jì)算總量；
模型蒸餾：訓(xùn)練一個(gè)較小的模型，該模型能夠復(fù)制原始模型的行為。
我們來看下 BERT 在這些策略上可以做什么。雷鋒網(wǎng)
1、架構(gòu)和優(yōu)化改進(jìn)
大規(guī)模分布式訓(xùn)練
加速 BERT 訓(xùn)練能夠想到的第一件事（甚至是第零件事）是將其分發(fā)到更大的集群當(dāng)中。雖然最初的 BERT 已經(jīng)使用了多臺(tái)機(jī)器進(jìn)行訓(xùn)練，但還有更多針對(duì) BERT 分布式訓(xùn)練的優(yōu)化解決方案，例如阿里巴巴 [1] 和英偉達(dá) [2] 的一些工作。
英偉達(dá)最近使用 NVIDIA DGX SuperPOD（具有 92 個(gè) DGX-2H 節(jié)點(diǎn)，共有 1472 個(gè) V100 GPU，理論上可以提供 190PFLOPS）刷新了 BERT 訓(xùn)練的記錄，在 53 分鐘內(nèi)訓(xùn)練出了 BERT-large 模型 [3]。當(dāng)然，這個(gè)訓(xùn)練也是非常昂貴的，除了英偉達(dá)，其他公司和個(gè)人很難去做這樣的事情。
還有另外一種更為巧妙優(yōu)化方案，這是一種被稱為 LAMB 的新的分層自適應(yīng)大批量優(yōu)化技術(shù) [4]，這種技術(shù)可以將 TPUv3 Pod 上的 BERT 訓(xùn)練時(shí)間從 3 天降低到 76 分鐘（1024 個(gè) TPUv3 芯片，可以提供超過 100PFLOPS），顯然，訓(xùn)練這個(gè)模型也不會(huì)便宜。
架構(gòu)
現(xiàn)在考慮更加架構(gòu)性的以及更少硬件的解決方案。
有一種基于對(duì)自注意層行為的觀察來訓(xùn)練 BERT 的漸進(jìn)式堆疊方法 [5]，該方法的分布局部地集中在其位置和 start-of-sentence token 周圍，且在淺層模型中的注意力分布與深層模型相似。為此，作者提出了將知識(shí)從淺層模型轉(zhuǎn)移到深層模型的堆疊算法，并逐步應(yīng)用堆棧來加快 BERT 訓(xùn)練的速度。憑借這種方式，作者獲得模型的訓(xùn)練時(shí)間相比原始 BERT 的訓(xùn)練時(shí)間縮短了約 25%，歸其原因則在于對(duì)于相同數(shù)量的步驟，訓(xùn)練一個(gè)小的模型需要的計(jì)算也相對(duì)較少。
還有一種方法是注意力矩陣的稀疏分解（也稱 Sparse Transformer，這項(xiàng)研究是由 OpenAI 完成的)[6] 和 block attention[7]，它們都屬于能夠減少內(nèi)存和計(jì)算總量的體系結(jié)構(gòu)改進(jìn)。
ALBERT
最后不得不提的是，已經(jīng)提交給 ICLR 2020 的一個(gè) BERT 的后代，被稱為 ALBERT（A Lite BERT）[8]。
ALBERT 結(jié)合了兩種參數(shù)降低的技術(shù)。
第一種是分解式嵌入?yún)?shù)化，將隱藏層的大小與詞嵌入的大小分開。這種分隔使得在不顯著增加詞匯表嵌入?yún)?shù)大小的情況下能夠更容易增加隱藏層的大小。
第二種是跨層參數(shù)共享。這種技術(shù)可以防止參數(shù)隨著網(wǎng)絡(luò)深度的增加而增大。
這兩種技術(shù)可以顯著降低 BERT 的參數(shù)數(shù)量，而不會(huì)嚴(yán)重影響性能，從而提高參數(shù)效率。
與 BERT-large 具有相似配置的 ALBERT，其參數(shù)能夠降低 18 倍，訓(xùn)練速度提高了月 1.7 倍。
相比嚴(yán)格調(diào)校的 RoBERTa，ALBERT 則同樣更勝一籌 [9]。
2、量化和修剪
量化 會(huì)降低模型權(quán)重的數(shù)值精度。通常情況下，使用 FP32（32 位浮點(diǎn)）來訓(xùn)練模型，然后可以將其量化為 FP16（16 位浮點(diǎn)），INT8（8 位整數(shù)），甚至將其量化為 INT4 或 INT1。于是模型尺寸將隨之減小 2 倍、4 倍、8 倍或 32 倍。這稱為 post-training quantization。
另一個(gè)選項(xiàng)是量化感知訓(xùn)練（也是較為困難和較為不成熟的）。這種方法的 FP16 訓(xùn)練現(xiàn)在已成為一種普遍的方法。而在 ICLR 2020 提交的論文中有一個(gè)比較有趣的工作 [10]，它使用 ResNet、GNMT、Transformer 的 8 位浮點(diǎn)表示獲得了最先進(jìn)的訓(xùn)練結(jié)果。
修剪 即從模型中刪除一些不重要或不太重要的權(quán)重（有時(shí)會(huì)是神經(jīng)元），從而產(chǎn)生稀疏的權(quán)重矩陣（或較小的圖層）。甚至還有一些研究直接去除掉與 Transformer 的注意力頭相對(duì)應(yīng)的整個(gè)矩陣。
量化，可以使用 TensorFlow Lite（用于設(shè)備上推斷的 TensorFlow 的部分）[11] 來執(zhí)行。TensorFlow Lite 提供了在移動(dòng)設(shè)備、嵌入式設(shè)備和 IoT 設(shè)備上轉(zhuǎn)換和運(yùn)行 TensorFlow 模型的工具，它支持訓(xùn)練后量化和量化感知訓(xùn)練。
另一個(gè)選擇是使用英偉達(dá)的 TensorRT 框架 [12]。英偉達(dá) TensorRT 是用于高性能深度學(xué)習(xí)推理的平臺(tái)，它包括深度學(xué)習(xí)推理優(yōu)化器，并且在運(yùn)行時(shí)能夠?yàn)樯疃葘W(xué)習(xí)推理應(yīng)用程序提供低延遲和高吞吐量。
英偉達(dá)最近發(fā)布了具有最新優(yōu)化功能的 TensorRT 6[13]，它可以在 T4 GPU 上僅用 5.8 毫秒對(duì) BERT-large 進(jìn)行推理，而在 V100 上甚至只需要 4.2 毫秒，對(duì)于 Titan RTX，速度可能會(huì)更快。
當(dāng)批處理數(shù)為 1 時(shí)，對(duì)于 340M 個(gè)參數(shù)的 BERT-large 模型，僅需 5.84 毫秒；對(duì)于 110M 參數(shù)的 BERT-Base 模型，則只需 2.07 毫秒。若將批處理數(shù)設(shè)為 128，你在使用 BERT-large 時(shí)甚至可以達(dá)到 250 個(gè)句子/秒的處理速度。
更多的數(shù)據(jù)我們這里就不再一一列舉了。
PyTorch 最近也宣布了在其 1.3 版本 [14] 中支持量化。盡管目前它還處于實(shí)驗(yàn)階段，但我們已經(jīng)可以使用它了，因?yàn)樵谄浣坛讨刑岬剿麄円呀?jīng)能夠?qū)?dòng)態(tài)量化應(yīng)用于將模型權(quán)重轉(zhuǎn)換為 INT8 的 LSTM 語言模型 [15]。
此外，還有一個(gè)眾所周知的 BERT 量化，稱為 Q-BERT。這項(xiàng)工作來源于「Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT」[16]。在這個(gè)研究中，作者甚至降低到 2-bit 的超低精度量化，但其性能相比于基線卻沒有顯著下降（僅下降 2.3%），而對(duì)應(yīng)的模型參數(shù)壓縮率最高可以達(dá) 13 倍，嵌入表壓縮率和激活的最高都為 4 倍。
3、蒸餾
另一個(gè)有趣的模型壓縮方法是蒸餾，這是一種將大型「teacher」網(wǎng)絡(luò)的知識(shí)轉(zhuǎn)移到較小的「student」網(wǎng)絡(luò)的技術(shù)，訓(xùn)練學(xué)生網(wǎng)絡(luò)來模仿教師網(wǎng)絡(luò)的行為。
Rich Caruana 及其合作者率先采用了這種策略。在他們先驅(qū)性的論文 [17] 中，他們提供了令人信服的證明：大型集成模型所獲得的知識(shí)可以轉(zhuǎn)移到單個(gè)小型的模型中。
Geoffrey Hinton 等人在他們的「Distilling the Knowledge in a Neural Network」{18} 論文中證明了這種技術(shù)可以應(yīng)用于神經(jīng)網(wǎng)絡(luò)模型。
DistilBERT
從 Hinton 開始，蒸餾的方法逐漸被應(yīng)用到了不同的神經(jīng)網(wǎng)絡(luò)中，當(dāng)然你可能也聽過 HuggingFace 提出的 DistilBERT，這是一種對(duì) BERT 的蒸餾。這項(xiàng)工作出自論文「DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter」[19]，目前已經(jīng)提交到 NeurIPS 2019。
DistilBERT 是一種較小的語言模型，受 BERT 的監(jiān)督而訓(xùn)練。在該模型中，作者刪除了令牌類型嵌入和合并器（用于下一個(gè)句子分類任務(wù)），并保持體系架構(gòu)其余部分不變，同時(shí)將層數(shù)減少了兩倍。
您可以在 HuggingFace（以前叫做 pytorch-transformers 和 pytorch-pretrained-bert）的 translators python 軟件包的幫助下，使用現(xiàn)成的 DistilBERT。該軟件包的 2.0.0 版本支持 TensorFlow 2.0 / PyTorch 互操作。
DistilBERT 的作者還使用了最新 RoBERTa 論文中的一些訓(xùn)練技巧，這些技巧表明 BERT 的訓(xùn)練方式對(duì)其最終性能至關(guān)重要。
DistilBERT 與 BERT 相比具有出人意料的結(jié)果：作者保留了 95％以上的性能，但參數(shù)卻減少了 40％。
比較 GLUE 基準(zhǔn)測(cè)試的開發(fā)集
在推斷時(shí)間方面，DistilBERT 比 BERT 快了 60％以上，比 ELMo + BiLSTM 快 120％。
推理速度
TinyBERT
就在幾天前，出現(xiàn)了一種新的 BERT 蒸餾方法，來自華為諾亞方舟實(shí)驗(yàn)室的劉群團(tuán)隊(duì)提出了 TinyBERT[20]。
為了構(gòu)建一個(gè)具有競(jìng)爭(zhēng)力的 TinyBERT，作者首先提出了一種新的 Transformer 蒸餾方法，來蒸餾嵌入 BERT 的知識(shí)。
具體來說就是，他們?cè)O(shè)計(jì)了幾個(gè)損失函數(shù)來適合 BERT 層的不同表示形式：
1、嵌入層的輸出；
2、Transformer 層派生的隱藏狀態(tài)和注意力矩陣;
3、預(yù)測(cè)層輸出的 logits 輸出。
論文中基于注意力的擬合則得益于最近的發(fā)現(xiàn)，即 BERT 學(xué)習(xí)的注意力權(quán)重可以捕獲大量的語言知識(shí)，這意味著語言知識(shí)可以很好地從教師 BERT 轉(zhuǎn)移到學(xué)生 TinyBERT。而在 BERT 的現(xiàn)有知識(shí)蒸餾的方法（如 Distilled BiLSTM_SOFT，BERT-PKD 和 DistilBERT）中卻忽略了這一點(diǎn)。
在這項(xiàng)工作中，作者還提出了一種新穎的兩階段學(xué)習(xí)框架，包括通用蒸餾和特定任務(wù)蒸餾。在通用蒸餾階段，未經(jīng)微調(diào)的原始 BERT 充當(dāng)教師模型，學(xué)生 TinyBERT 通過在通用領(lǐng)域?qū)Υ笮驼Z料庫執(zhí)行通常的 Transformer 蒸餾來學(xué)習(xí)模仿教師的行為。他們獲得了可以對(duì)各種下游任務(wù)進(jìn)行微調(diào)的通用 TinyBERT。在特定任務(wù)蒸餾階段，他們將數(shù)據(jù)進(jìn)行擴(kuò)充，來提供更多與任務(wù)相關(guān)的材料供教師-學(xué)生學(xué)習(xí)，然后在增強(qiáng)的數(shù)據(jù)上重新執(zhí)行 Tranformer 蒸餾。
這個(gè)兩階段方法對(duì)于提升 TinyBERT 的性能和泛化能力是必不可少的。
TinyBERY 在實(shí)驗(yàn)上取得了非常的成績，相對(duì)于 GLUE 數(shù)據(jù)集的 BERT-base，其性能并沒有下降多少，而推理參數(shù)小了 7.5 倍，推理時(shí)間快了 9.4 倍。
TinyBERT 與其他基準(zhǔn)的比較
我們期待他們能夠?qū)⑦@種方法應(yīng)用到 BERT-large 和 XLNet 等大型模型中，同樣也期待他們開放源碼。
其他蒸餾方法
除了 DistilBERT 和 TinyBERT 外，還有其他一些為大家所熟知的蒸餾方法。
（2019/03）「Distilling Task-Specific Knowledge from BERT into Simple Neural Networks」[21]
這篇論文將 BERT 蒸餾到單層 BiLSTM 中，取得了與 ELMo 可比的結(jié)果，同時(shí)使用的參數(shù)減少了大約 100 倍，推理時(shí)間減少了 15 倍。
BiLSTM_SOF 是 TBiLSTM 的蒸餾，后者是在 soft logit 目標(biāo)上訓(xùn)練出來的。
（2019/08）「Patient Knowledge Distillation for BERT Model Compression」[22]
這篇論文提出了一種耐心知識(shí)蒸餾的方法，這是首次嘗試使用教師的隱藏狀態(tài)，而不僅僅是最后一層的輸出。他們的學(xué)生模型從教師模型的多個(gè)中間層「耐心」地學(xué)習(xí)來獲得更多知識(shí)。在他們的耐心蒸餾知識(shí)框架中，只訓(xùn)練學(xué)生模仿中間層的 [CLS] 令牌的表示形式。代碼已公開 [23]。
（2019/09）「Extreme Language Model Compression with Optimal Subwords and Shared Projections」
這是最近提交到 ICLR 2020 的一篇論文，這篇論文專注于一種用于訓(xùn)練詞匯量顯著較小、嵌入和隱藏狀態(tài)維度較低的學(xué)生模型的知識(shí)蒸餾技術(shù)。作者采用了雙重訓(xùn)練機(jī)制，可以同時(shí)訓(xùn)練教師和學(xué)生模型，從而能夠獲得針對(duì)學(xué)生詞匯的最佳詞嵌入。該方法能夠?qū)?BERT-base 模型壓縮 60 倍以上，而下游任務(wù)指標(biāo)只有很小的下降，從而使得語言模型占用的空間只有不到 7MB。
TinyBERT 的結(jié)果似乎更好，但一個(gè) 7MB 的類 BERT 模型簡(jiǎn)直爽的不能再爽！
需要強(qiáng)調(diào)，以上所介紹的方法并不互相沖突，所以我們期待能夠?qū)⑦@些方法協(xié)同起來使用來加速 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
雷鋒網(wǎng)(公眾號(hào)：雷鋒網(wǎng))報(bào)道。
