?作者 | 徐超

單位 | 微軟亞洲互聯(lián)網(wǎng)工程院

研究方向 | 文本相關(guān)性、多語言擴展

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

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

裁剪

BERT 模型包含一個 embedding 層與多個 Transformer 層。如果你的語料比較小，任務(wù)比較簡單，可以對 BERT 模型做適當?shù)牟眉?#xff1a;?

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

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

知識蒸餾

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

2.1 Embedding層的壓縮

2.1.1?維度縮減

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

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

2.1.2?Factorized embedding parameterization

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

2.2?Transformer層的壓縮

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

2.2.1 Cross-layer parameter sharing?

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

2.2.2 寬度壓縮?

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

2.3 蒸餾方法

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

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

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

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

2.4?Task Specific蒸餾

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

2.4.1?完整蒸餾流程

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

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

2. 收集具體任務(wù)相關(guān)數(shù)據(jù)集，然后用 BERT 模型在這個數(shù)據(jù)集上做微調(diào)得到 Teacher Model；

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

Task Specific 的損失函數(shù)一般采用：

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

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

2.4.2?Capacity Gap

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

2.4.3?Spherical Knowledge Disitllation (SKD)

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

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

2.4.4?Multi-Stage Distillation

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

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

量化

量化是一種信息的有損壓縮方式，在神經(jīng)網(wǎng)絡(luò)里面的應(yīng)用是將高精度的浮點數(shù)運算轉(zhuǎn)化成低精度的整數(shù)運算，以加速模型的而推理計算。

量化方式有兩種：

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

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

3.1?對稱線性量化

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

其中 是輸入數(shù)據(jù) x 的尺度參數(shù)，M 是當量化位數(shù)為 b 時的最大量化結(jié)果，即：

在訓(xùn)練過程中 weights 的尺度參數(shù)可以設(shè)為：

activations 的尺度參數(shù)可基于下式，其中 EMA 為指數(shù)加權(quán)平均。

3.2?Integer Only BERT

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

3.3?改進量化

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

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

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

3.4 BinaryBERT

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

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

其中：

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

其他方法

4.1 計算流圖優(yōu)化?

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

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

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

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

4.2?硬件加速?

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

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

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

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

4.3?FPGA硬件加速

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

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

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

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

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

總結(jié)

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

特別鳴謝

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

參考文獻

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.

更多閱讀

#投 稿?通 道#

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

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

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

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優(yōu)質(zhì)內(nèi)容，可以是最新論文解讀，也可以是學(xué)術(shù)熱點剖析、科研心得或競賽經(jīng)驗講解等。我們的目的只有一個，讓知識真正流動起來。

📝?稿件基本要求：

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

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

? PaperWeekly 尊重原作者署名權(quán)，并將為每篇被采納的原創(chuàng)首發(fā)稿件，提供業(yè)內(nèi)具有競爭力稿酬，具體依據(jù)文章閱讀量和文章質(zhì)量階梯制結(jié)算

📬?投稿通道：

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

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

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

△長按添加PaperWeekly小編

🔍

現(xiàn)在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關(guān)注」訂閱我們的專欄吧

·

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。