?作者 | 徐超

單位 | 微軟亞洲互聯網工程院

研究方向 | 文本相關性、多語言擴展

基于 Transformer 的預訓練模型，尤其是 BERT，給各種 NLP 任務的 performance 帶來了質的飛躍。如今 pretrained model + Fine tune 幾乎已經成為 NLP modeling 工作的標準范式。然而這些模型卻是越來越重，如 RoBERTa-large 有 3.55 億參數，GPT2-xl 有 15 億參數，GPT3 的參數達到了 1750?億！

在有著三高問題（特點？）的互聯網系統中使用這種模型面臨著很大的性能瓶頸，甚至已經很難將一個 SOTA 的模型直接搬到線上進行實時的推理計算。本文將介紹最廣為使用的預訓練模型 BERT 的推理加速技術，這些技術可以應用到其他類似 Transformer Based 預訓練模型中。

裁剪

BERT 模型包含一個 embedding 層與多個 Transformer 層。如果你的語料比較小，任務比較簡單，可以對 BERT 模型做適當的裁剪：?

Embedding 層的裁剪：如果用的是 Multilingual BERT 可以將用不到的語言的 Word Embedding 裁剪掉。可以用 BERT 的 vocabulary 對你的 target 語料進行 BPE 切詞，并統計詞頻，對于頻次特別低的詞的 Embedding 也可以裁剪掉。?

Transformer 層的裁剪：BERT 中包含了 12 層的 Transformer，大的有 24 層，其他模型甚至更多，但實際你的任務如果不是很復雜，可能未必需要這么多層。由于底層的 Transformer 善于捕捉細粒度的信息，一般有限選擇底層的 Transformer 層。常見的選擇方式是選擇底部幾層，如第 1、2、3 層，或者間隔選，如第 1、3、5 層。

知識蒸餾

原始的 BERT 模型是過參數化的，存在很大的壓縮空間。知識蒸餾應該是主要的模型壓縮技術，他是將 BERT 模型（Teacher Model）里面的知識通過特殊的訓練方式濃縮到一個較小的模型（Student Model）里面，這個較小的模型 performance 并沒有顯著下降甚至更好。在設計這個 Student Model 時也是從 Embedding 與 Transformer 層兩個角度對 BERT Model 進行了壓縮。

2.1 Embedding層的壓縮

2.1.1?維度縮減

這是最簡單粗暴的方式，BERT 模型的 Embedding 維度有 768 與 1024 兩個版本，你在設計 Student Model 時可以用 512 甚至更低的維度。這時候如果想復用 BERT 的 Transformer 的參數，那么需要對 Embedding 做一個線性變換，將其轉換成 768 或者 1024 維。可通過添加一個全連接層來實現。

當然也可以將這個線性變換用更復雜的 MLP 網絡代替，常見變換如下，其中 DeFINE 模型里面使用的 HGT 取得了 SOTA 的效果，比一般的矩陣分解技術都好。

2.1.2?Factorized embedding parameterization

這是 ALBERT 里面提出的方式，將 embedding matrix 分解為兩個大小分別為 和 矩陣，也就是說先將單詞投影到一個低維的embedding空間 ，再將其投影到高維的隱藏空間 。這使得 embedding matrix 的維度從 減小到 。當 時，參數量減少非常明顯。在實現時，隨機初始化 和 的矩陣，計算某個單詞的表示需用一個單詞的 one-hot 向量乘以 維的矩陣（也就是 lookup），再用得到的結果乘 維的矩陣即可。

2.2?Transformer層的壓縮

BERT 的 Transformer 層既深又寬，所以可以從深度與寬度兩個維度進行壓縮。眾所周知，深度學習的主要威力就在深，所以直接減少 Transformer 結構的層數，會降低模型的非線性擬合能力，對 performance 產生很大影響，一個可行的方法是層與層之間的參數共享，即 Cross-layer parameter sharing。

2.2.1 Cross-layer parameter sharing?

這種方法是在 ALBERT 中使用的，即多個層使用相同的參數，它不能較少推理的時間但可以減少模型的體積。參數共享有三種方式：只共享 feed-forward network 的參數、只共享 attention 的參數、共享全部參數。?

2.2.2 寬度壓縮?

另一個更優雅的方式是壓縮寬度，使用窄深的模型，下圖是 MobileBERT 用的模型架構。最終產出的是窄深的 MobileBERT，但需先訓練一個比較寬的 model，然后將其知識蒸餾到 MobileBERT 中。其中寬模型與窄模型使用了相同的 Embedding 層，這個 embedding 層通過線性變換來適配各自的 Transformer 層。

2.3 蒸餾方法

BERT 的蒸餾方法已經有過很多研究，如 BERT-PKD、DistilBERT 以及 TinyBERT 等，這些方法的賣點主要在 loss 的設計上。這里介紹下 TinyBERT 的 Loss 設計，應該說考慮的比較全面主要包含三大塊：Attention Loss，Hidden Loss 與 Embedding Loss。

Attention Loss：Teacher 與 Student 的每層的注意力向量 A 的均方誤差，其中 h 是 attention 的頭數。

Hidden Loss：Transformer 層的輸出的均方誤差，由于輸出維度不一致，需要使用 進行維度轉換：

Embedding Loss：同樣使用均方誤差定義，并使用 進行維度轉換：

2.4?Task Specific蒸餾

前面的蒸餾方式是針對 BERT 模型本身的，學習的是 BERT 中通用的語言表征方式，并未涉及到具體的任務，這種蒸餾被稱為 Task Agnostic 蒸餾。預訓練模型只是提供了一個初始參數，針對具體的任務，還要在相應的數據集上做微調，得到我們具體任務的 model，然后我們對這個針對具體 NLP 任務的 model 做蒸餾，也就是 Task Sepcific 蒸餾。

2.4.1?完整蒸餾流程

以 TinyBERT 作為 Student Model 為例， 整個蒸餾的流程為：

1. 用上一小節介紹的 Task Agnostic 方法將 BERT 中的知識蒸餾到 TinyBERT 中；?

2. 收集具體任務相關數據集，然后用 BERT 模型在這個數據集上做微調得到 Teacher Model；

3. 以第一步的 TinyBERT 的參數來初始化 Student Model，用 Task Specific 蒸餾方法將 Teacher Model 的知識蒸餾到 Student Model。

Task Specific 的損失函數一般采用：

Teacher Model 輸出的 logits 與 Student Model 輸出的 logits 之間的均方誤差：

對 Student Model 的 logits 升溫后的交叉熵損失：采用一個廣義的 softmax 函數，T 為溫度，值越大，輸出的概率分布越平緩，越小則分布越陡峭，當 T=0?時就變成了 one-hot 編碼。在做 inference 時 T=1，在做訓練時為大于 1 的某個數。

2.4.2?Capacity Gap

在知識蒸餾問題中，并非 Teacher Model 越好 Student Model 就越好。當 Teacher Model 比 Student Model 大很多時，Capacity Gap 問題將會變得非常嚴重。下圖中的 Student Model 都是 ResNet14，當 Teacher Model 的體量超過 ResNet32 時，蒸餾的結果卻出現了明顯的 drop。同時 KD Loss 隨著 Teacher Model 的增大而不斷增加，說明 Student Model 對 Teacher Model 的模仿能力或者擬合程度越來越低。

2.4.3?Spherical Knowledge Disitllation (SKD)

一般隨著模型的增大，會對預測的 label 更加 confident，即分配更高的概率。如下圖左側，Teacher Model 對預測的 label 分配的概率時 0.984，而 Student Model 分配的概率是 0.73。而常規的知識蒸餾使用了 Soft Target，即讓 Student Model 去學習 Teacher Model 的輸出概率分布，然而 Student Model 在體量上的劣勢就注定了它很難在這方面向 Teacher Model 看齊。

一個解決方案是，鈍化 Teacher Model 的輸出概率，使得 Student Model 容易模仿。具體做法是對模型輸出的 logits 進行尺度變換，成為一個單位向量 ，然后乘以 logits 的均值 ，使得變換后的的 logits 跟原始 logits 維持在同一個量級，這么做是為了盡量減少對有監督學習的影響。最終的 SKD Loss 為：

2.4.4?Multi-Stage Distillation

另一個緩解 Capacity Gap 的策略是使用多階段的蒸餾，一般使用兩階段，引入一個 size 適中的網絡，即 Teacher Assistant，作為 Teacher 與 Student 之間的過渡。Teacher 與 TA 以及 TA 與 Student 之間的 Capacity Gap 都不是很嚴重，先將 Teacher 的知識蒸餾到 TA，再將 TA 的知識蒸餾到 Student，這就避免了 Gap 過大時蒸餾效果出現劇烈 drop 的情況。

那么 TA 的 size 選多少合適呢？一般并不是選擇兩者體積的平均，而是在各自單獨訓練的情形下，性能介于 Teacher 與 Student 之間的網絡的 size。例如，在下圖重藍色的折線表示隨著模型層數的增加，模型的準確率，如果 Student 是兩層，Teacher 是 10 層，那么 TA 應該選 4 層，因為它的性能接近前兩者的平均值。

量化

量化是一種信息的有損壓縮方式，在神經網絡里面的應用是將高精度的浮點數運算轉化成低精度的整數運算，以加速模型的而推理計算。

量化方式有兩種：

訓練后量化：模型訓練好之后再進行量化，這種方法一般對大模型效果不錯，對小模型則效果比較差；

訓練感知量化：在訓練過程中對模型參數量化，inference 與 training 使用的是量化過后的 weights 和 activation。

3.1?對稱線性量化

為了保證精度，一般都是用訓練感知量化。在 BERT 中經常采用的量化算法是對稱線性量化（symmetric linear quantization）。

其中 是輸入數據 x 的尺度參數，M 是當量化位數為 b 時的最大量化結果，即：

在訓練過程中 weights 的尺度參數可以設為：

activations 的尺度參數可基于下式，其中 EMA 為指數加權平均。

3.2?Integer Only BERT

這種量化的目的是是的整個推理過程中只有整數運算，包括量化部分與激活函數的計算，這樣一來就可以使用 Turning Tensor Cores 或者傳統的 integer-only ARM 硬件來獲得進一步的加速效果（大約 2.4~4 倍）。

3.3?改進量化

用一個固定位數（如最常見的 8 位）的整數對所有參數進行量化是有點粗暴的，可以從兩個方面進行改進：

混合精度量化：不同的參數對量化誤差的敏感程度不一樣，所以使用不同的位數進行量化，對于 Activation 可以選用更高的位數，如 16 位，而對于 Embedding 可以選用更低的位數，如 2~4 位。

分組量化：增加量化的粒度可以提升效果但會增加計算成本，一個折衷的方式是對參數進行分組，組內的參數共享量化參數（如尺度參數與位置參數），最常見的是對 Embedding 的維度進行分組。

3.4 BinaryBERT

量化的極限是 BinaryBERT，模型的 Weights 與 Embeddings 都用 1 位的整數表示，即要么是 0 要么是 1，模型的體量可以縮減為原來的 1/24，推理速度提升一個數量級。

但 BinaryBERT 難以直接訓練，一個折中的做法是先訓練一個 ternary BERT（量化位數 2bit），然后采用一種被稱為 ternary weight splitting (TWS) 的操作獲得 BinaryBERT 的權重，最后對這個 BinaryBERT 做微調。TWS 操作如下：

其中：

有點類似于用兩個 Binary 數據來表示一個 Ternary 數據，所以 BinaryBERT 的 size 也是 ternary BERT 的兩倍。

其他方法

4.1 計算流圖優化?

神經網絡是一個計算流圖，原始的執行流程還可以還很多優化，比如算子的合并，運行時不需要部分的裁剪等。常見的針對神經網絡計算流圖優化的引擎主要有 ONNX Runtime 與 TVM，其中 ONNX Runtime 使用最廣。

ONNX Runtime 是由微軟推出的推理優化框架，支持多種運行后端包括 CPU，GPU，TensorRT，DML 等，并且對通用的 ONNX 格式模型提供了原生的支持。其基本原理是先執行一遍模型的推理，然后自動將模型中用不到的部分修剪掉，并將部分算子做了合并，還能對網絡的拓撲結構做優化，最終生成一個精簡的靜態計算流圖。

使用起來也非常簡單，只要加載模型然后執行一次推理運算即可，具體詳見官方文檔：

https://onnxruntime.ai/docs/api/python/api_summary.html

4.2?硬件加速?

前面介紹的推理加速方法都是對模型本身的優化，換個角度思考，也可以從硬件方面去進行加速，不過由于硬件門檻高，這方面的案例并不多，不過大廠好像都在這塊有投入。

神經網絡處理器：由于芯片定制化生產技術的成熟，谷歌在 2015 年就推出了專門用于神經網絡計算的 TPU，如今 AI 應用潤物細無聲，占據了大部分的算力消耗，各大硬件廠商都在發力神將網絡處理器（NPU）。

ASIC：在某些人眼中，可能任何處理器都不是最優的只有 ASIC 才是最優的，也就是為某個應用定制專業的芯片。這塊，如果模型已經高度產品化了，如語言識別、圖像識別、機器翻譯等可以考慮使用。

FPGA：這塊也炒的比較火，微軟已經在數據中心使用了大量的 FPGA 代替 CPU 做科學計算。單從算力本身來說 FPGA 的主頻遠比不上 CPU 與 GPU，但他的優勢在于并行度，對于大的矩陣計算，CPU 需要很多時鐘周期才能完成，GPU 拓展了帶寬比 CPU 有了明顯的優勢，但 FPGA 在并行上有無與倫比的優勢，只要門電路足夠，可以根據需要不斷增加并行度，并且它完成一件事情的時間是確定的，每個并發完成的時間完全一致，不存在時間片、線程與資源沖突問題。FPGA 的另一個優勢是完成同樣的計算功耗比 CPU 低很多，GPU 就更不用說了。但這一塊實際使用起來并沒有那么容易，可以應用的場景非常少，下一節做個簡單的案例介紹。

4.3?FPGA硬件加速

FPGA 本身的主頻非常低，跟 CPU 差遠了，而且存儲資源非常小，數據吞吐能力差，完全用 FPGA 實現一個 BERT 模型理論上不是不行，但肯定是土豪的玩具。一個折衷的辦法是與 CPU 協同工作，把 CPU 不擅長且 memory cost 較低的矩陣計算任務交給它，也就是 BERT 中的 Encoder 層，CPU 主要處理 Embedding 層與 Task Specific 層。當然，由于 FPGA 的 memory 限制，需要把 BERT 模型量化成 Integer Only BERT，整個系統的架構如下：

CPU 計算的中間結果與模型的權重都存放在 off-chip memory 上，通過 AXI4 接口傳送給 FPGA。FPGA 內部實現了三個計算單元用于加速計算：

1. PE (Processing Element)：用于不同位寬的矩陣乘法計算。

2. Softmax Core：用于進行 Softmax 計算。主要是將 的結果轉化為 Attention 權重，但指數計算非常消耗資源。好在 Softmax 函數具有平移不變性，而所有數據都是整數，所以使用了一個 Lookup Table 來代替指數變換。

3. LN Core：用于 Layer Normalization 計算，由于 LN 運算中很多 Element Wise 的加乘運算與非線性變換，單獨依靠 PE 無法完成，所以針對這部分做了專門的設計。 整個數據流圖如下，FPGA 從 off-chip memory buffer 中讀取數據，為了減少 FPGA 的數據吞吐量，將每個階段分成幾個子階段，每個子階段只加載需要的數據，最后將計算結果寫入 goff-chip memory buffer 供 CPU 做 Task Specific 層的運算。

總結

面這些技術未必全部都需要做，一般經過蒸餾與 ONNX Runtime 的推理優化，可以將模型從 1~2G 壓縮到 30~60M，推理速度也至少提升一個數量級，基本滿足線上需求，模型的 performance 的 drop 一般在 5% 以內。量化雖然能顯著縮小模型的體積，但對小模型 ROI 較低，一般蒸餾與量化兩種方法二選一就行了，很少同時做。FPGA 本身算力較差，很難提升推理速度，與 CPU/GPU 相比，主要優勢在于更高的并行計算能力與更低的功耗比，但高度定制化，使用門檻高，設計的不好很可能會弄巧成拙，主要用在資源受限的邊緣計算設備上。

特別鳴謝

感謝 TCCI 天橋腦科學研究院對于 PaperWeekly 的支持。TCCI 關注大腦探知、大腦功能和大腦健康。

參考文獻

V. Sanh, L. Debut, J. Chaumond, and T. Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter, 2020.?

S. Sun, Y. Cheng, Z. Gan, and J. Liu. Patient knowledge distillation for bert model compression, 2019.?

Sachin Mehta, Rik Koncel-Kedziorski, Mohammad Rastegari, Hannaneh Hajishirzi. DeFINE: DEEP FACTORIZED INPUT TOKEN EMBEDDINGS FOR NEURAL SEQUENCE MODELING, 2020.?

Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, Denny Zhou. MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices, 2020.?

Jia Guo, Minghao Chen, Yao Hu, Chen Zhu, Xiaofei He, Deng Cai. Reducing the Teacher-Student Gap via Spherical Knowledge Disitllation.?

Seyed-Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Hassan Ghasemzadeh. Improved Knowledge Distillation via Teacher Assistant: Bridging the Gap Between Student and Teacher, 2019.?

Ofir Zafrir, Guy Boudoukh, Peter Izsak, Moshe Wasserblat. Q8BERT: Quantized 8Bit BERT, 2019.?

Sehoon Kim, Amir Gholami, Zhewei Yao, Michael W. Mahoney, Kurt Keutzer. I-BERT: Integer-only BERT Quantization, 2021.?

Yelysei Bondarenko, Markus Nagel, Tijmen Blankevoort. Understanding and Overcoming the Challenges of Efficient Transformer Quantization. 2021.?

Zejian Liu , Gang Li , Jian Cheng. Hardware Acceleration of Fully Quantized BERT for Efficient Natural Language Processing, 2021.

更多閱讀

#投 稿?通 道#

?讓你的文字被更多人看到?

如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。

📝?稿件基本要求：

? 文章確系個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標注?

? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發送，要求圖片清晰，無版權問題

? PaperWeekly 尊重原作者署名權，并將為每篇被采納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算

📬?投稿通道：

? 投稿郵箱：hr@paperweekly.site?

? 來稿請備注即時聯系方式（微信），以便我們在稿件選用的第一時間聯系作者

? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿

△長按添加PaperWeekly小編

🔍

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

·

總結

以上是生活随笔為你收集整理的预训练模型参数量越来越大？这里有你需要的BERT推理加速技术指南的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。