浅谈BERT/Transformer模型的压缩与优化加速
?作者 | 姚益武
單位 | 阿里巴巴集團
研究方向 | AI算法與工程架構
前言
BERT/Transformer?結構及其變體，已成為自然語言處理（NLP）、語音識別 （ASR）等領域的主流序列建模結構。并且，相比于卷積操作的局部特征提取能力、以及平移不變性，Self-Attention/Cross-Attention 的全局 Context 信息編碼或長時依賴建模能力，能夠使圖像/視覺模型具備更強的相關性建模能力（更好的全局視野）、與動態建模能力（領域或樣本適應性）。
因此在計算機視覺領域 （CV），Vision Transformer 模型結構（如 ViT、Deit、Swin-Transformer、PVT、SegFormer 與 DETR 等）也日益流行、方興未艾。然而，Transformer 模型的高計算復雜度與參數量，限制了模型的推理部署（包括服務端與移動端應用），其計算、存儲與運行時內存等資源開銷都成為限制因素（存在大量密集計算算子、與訪存算子，如 BatchMatMul、LayerNorm、Transpose/Reshape 等）。例如，從標準 Transformer layer 的 Tensor 處理結構來看，MHA 的 Context 特征計算與特征長度的平方、Embedding 特征維度成正比：
Standard Self-Attention (X=Y) / Cross-Attention (X!=Y):
?Standard FFN:
本文從以下幾個維度舉例說明，解析 Bert/Transformer 模型的加速方法，體現了從建模依賴、數據依賴到硬件依賴的優化層次：
模型結構精簡化與知識蒸餾
模型量化（Quantization）
網絡結構搜索（NAS: Network Architecture Search）
計算圖優化（計算圖等價變換）
推理優化引擎（Faster Transformer / Torch-TensorRT / AICompiler）
模型結構精簡化與知識蒸餾
2.1 Lite-Transformer
論文標題：
Lite Transformer with Long-Short Range Attention
收錄會議：
ICLR 2020
論文鏈接：
https://arxiv.org/abs/2004.11886
代碼鏈接：
https://github.com/mit-han-lab/lite-transformer
Lite Transformer 是一種融合卷積與 Self-Attention?操作的、高效精簡的 Transformer 結構，可應用于 NMT、ASR 等序列生成任務。其核心是長短距離注意力結構（LSRA：Long-Short Range Attention），能夠有效融合卷積與 Attention 機制。
LSRA 將輸入 Embedding 沿 Feature 維度 split 成兩部分，其中一部分通過 GLU（Gate Linear Unit）、一維卷積，用以提取局部 Context 信息；而另一部分依靠 Self-attention，完成全局相關性信息編碼。一維卷積的引入，能夠減輕計算量與參數量。Lite Transformer 核心結構如下，首先將 FFN 的寬度攤平（flatten），其次引入 LSRA 以替換 Self-Attention：
2.2 SAN-M
論文標題：
SAN-M: Memory Equipped Self-Attention for End-to-End Speech Recognition
收錄會議：
INTERSPEECH 2020
論文鏈接：
https://arxiv.org/abs/2006.01713
SAN-M 表示 Self-Attention 與 DFSMN 記憶單元的融合，是一種 Transformer ASR 模型。DFSMN 適合捕獲局部信息，Self-Attention 模塊具備較強的長時依賴建模能力，因此二者存在互補性。SAN-M 通過將兩個模塊的特性融合在一起，實現了優勢互補。Biasic Sub-layer 表示包含了 SAN-M 的 Self-Attention 模塊，DFSMN 添加在 values 后面，其輸出與Multi-head Attention（MHA）相加：
2.3 MiniLM（知識蒸餾）
論文標題：
MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers
論文鏈接：
https://arxiv.org/abs/2002.10957
代碼鏈接：
https://github.com/microsoft/unilm/tree/master/minilm
針對 NLP 任務，深度自注意力知識蒸餾（Deep Self-Attention Distillation），通過遷移 Teacher model 最后一層 Self-Attention 模塊的 Attention score 信息與 Value relation 信息，可有效實現 Student model 的誘導訓練。只遷移最后一層的知識，能夠直接遷移語義信息，顯得簡單有效、訓練速度更快；而且相比于層間特征遷移，不需要手動設計 Teacher-student 之間的層對應關系。Attention score 信息與 Value relation 信息的知識遷移如下：
Attention score transfer:
Value relation transfer:
若選擇的 Self-attention layer 為 LSRA，除了在 Multi-head Attention（MHA）分支遷移 Attention score 與 Value relation；在 CNN 分支需要遷移 Feature map 的信息，這里主要計算 AT loss：
式中 表示學生網絡里第 j 個網絡層的特征輸出， 表示 Teacher network 里第 j 個 group 的特征輸出。
2.4 Switch-Transformer
論文標題：
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
論文鏈接：
https://arxiv.org/abs/2101.03961
代碼鏈接：
https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/moe.py
針對內容理解與生成、以及多模態特征表征等 AI 任務，基于?MoE（Mixture of Experts）單元的大模型的參數規模不斷擴展（Switch-Transformer?是其中的典型代表之一），但大模型對算力的需求、被 MoE 的稀疏激活（Sparse activation）或動態路由（Dynamic routing）機制有效緩解。
MoE 單元屬于典型的 Sample 驅動型結構設計，使得 MoE-based 模型具備良好的動態建模能力（動態卷積、Self-Attention 等亦如是），因此 MoE 能夠非常好的適應多任務/多模態、甚至在線實時的場景應用需求。Switch-Transformer 的 MoE 單元結構如下，具體執行過程包括動態路由（Dynamic routing）、數據分派（Data dispatch）、專家計算（Expert computation）與結果合并（Result combine）：
模型量化（Quantization）
3.1 Transformer ASR 模型壓縮
論文標題：
Extremely Low Footprint End-to-End ASR System for Smart Device
收錄會議：
INTERSPEECH 2021
論文鏈接：
https://arxiv.org/abs/2104.05784
針對 Transformer ASR，文章提出了聯合隨機稀疏與 PTQ 量化（KL 量化、ADMM 與混合精度設置）的壓縮策略，整體實現了 10 倍壓縮，且絕對精度損失約 0.5%。總體流程如下：
模型稀疏化，更新 Weight 重要性，漸進式增加稀疏度：
KL 量化，計算 Activation 量化參數：
ADMM，優化 Weight 量化參數：
混合精度量化設置，減少量化誤差：
此外，針對 Transformer 模型的量化，需要討論全網絡 INT8 量化與計算的意義：
全網絡 INT8 量化：
同時減少計算密集算子與訪存密集算子的開銷；
實現模型壓縮，INT8 模型是 FP32 模型的 1/4；
Fully INT8 Attention Path：
基本的全 INT8 通路：
Dense layer 的輸入與輸出均為 INT8 Tensor；
Reshape 與 Transpose 的操作對象為 INT8 Tensor，節省內存開銷：
Dense+Reshape+Transpose、與 Transpose+Reshape+Dense 可以實現 Op fusion；
BatchMatMul、Softmax 的操作對象保留為 FP32 Tensor，確保模型預測精度；
Encoder的Self-Attention：
X=Y；
Decoder 的 Self-Attention：
X=Y，Batch_size=Batch_size * Beam_size，Seq_len=1；
k_dense 與 v_dense 的輸出會添加到 Cache，按照 INT8 類型搬運數據，可節省訪存開銷：
3.2 PTQ for Vision Transformer
論文標題：
Post-Training Quantization for Vision Transformer
論文鏈接：
https://arxiv.org/abs/2106.14156
代碼鏈接：
https://github.com/hahnyuan/PTQ4ViT
針對 Vision Transformer 的 PTQ 量化，分別針對 FFN 與 Self-Attention，提出了 Similarity-aware Quantization 與 Ranking-aware Quantization：
Similarity-aware Quantization：基于優化方式確定 Weight 與 Activation 量化的最優量化參數，并以 Pearson 相似度作為 Target loss；
Ranking-aware Quantization：在相似度 Loss 基礎之上，疊加 Attention map 的 Ranking loss 作為約束，以準確感知不同 Attention map 的重要性排序，確保 Self-attention 量化的有效性（優先保留重要 Attention map 的信息）；
3.3 PTQ4ViT
參考博客文章：
https://blog.csdn.net/nature553863/article/details/121564248
3.4 Q-Bert
論文標題：
Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT
收錄會議：
AAAI 2020
論文鏈接：
https://arxiv.org/abs/1909.05840
針對 NLP 任務（MNLI、SQuAD 等），Q-Bert 是首個實現 Bert 模型混合精度量化的 QAT 算法。Bert 模型是由 Embedding layer、Encoder layer 以及 Task-specific output layer 構成的前饋結構模型，且 Encoder layer 包含了 MHA 與 FFN 等子模塊。
不同的 Bert layer 體現了不同的量化敏感度（Quantization Sensitivity），配以不同的量化比特數，直接決定了量化后模型的預測精度。通常 Bert 的 Embedding layer 與 Output layer 的量化敏感度相對較高，故 Embedding layer 會配以較高的量化比特數（如 8bits），Output layer 則保留為 FP32 浮點實現。另外，盡管不同 Encoder layer 的結構相一致，但由于對 Context 信息編碼的貢獻有所區別，因此表現出不同的量化敏感性。
量化敏感度與混合精度量化：Q-Bert 采用線性均勻量化方式（min-max 量化），并通過 Dense layer 的 Hessian 信息（二階梯度）定義其量化敏感度：
式中 Hessian 矩陣特征值（Top eigenvalues）的均值，反映了針對同一個下游任務（相同訓練數據分布），不同 Dense layer 的量化敏感度；Hessian 矩陣特征值的標準差，反映了針對不同的下游任務（不同訓練數據分布），Dense layer 量化敏感度的占比。然后，基于不同 Dense layer 的量化敏感度排序，可確定 Bert 模型的混合精度量化配置，并進一步執行 QAT 微調訓練。
Group-wise量化：為了進一步提高量化后模型的精度，將 MHA 的 Q、K、V 與 O 權重矩陣按 Attention heads 分成不同的 Group，每個 Group 定義各自的量化范圍（min-max 范圍），從而更精細的保證不同 Group 的量化分辨率。而每個 Group 又可以根據實際需求、進一步劃分為不同的 Sub-group，以滿足更為苛刻的部署精度需求：?
網絡結構搜索（NAS）
4.1 AdaBert
論文標題：
AdaBERT: Task-Adaptive BERT Compression with Differentiable Neural Architecture Search
收錄會議：
IJCAI 2020
論文鏈接：
https://arxiv.org/abs/2001.04246
針對 NLP 任務，AdaBert?通過可微分搜索、與知識蒸餾，實現卷積類型的 NLP 模型搜索：
搜索空間：類似于 DARTS 的 Cell 設計，實現搜索空間的構造；區別于 Vanilla CNN，DARTS 的 Cell 單元結構體現出了非規整性，可能不利于實際推理引擎的加速；
搜索策略：基于 Gumbel Softmax 實現網絡結構的隨機采樣（類似于 FBNet）、與可微分搜索，并且采樣概率隨著訓練的進行、逐漸銳化，逼近 Argmax 采樣；另外，基于 FLOPS 與 Model size 構造了 Efficiency-aware Loss，作為搜索訓練的資源約束，實現 Hardware-aware 搜索；
知識蒸餾：多層次、任務相關的知識遷移；
4.2 DynaBert
論文標題：
DynaBERT: Dynamic BERT with Adaptive Width and Depth
收錄會議：
NeurIPS 2020 (Spotlight)
論文鏈接：
https://arxiv.org/abs/2004.04037
代碼鏈接：
https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/DynaBERT
針對 NLP 任務，DynaBert 模型壓縮策略可實現多個維度的壓縮搜索（主要是 Width 與 Depth 維度）。DynaBert 是一種動態可伸縮性質的 One-shot NAS：首先基于預訓練 Bert 模型（或其變體），進行多維度正則化訓練；正則化訓練之后，按照不同的資源約束需求，能夠進行網絡參數的裁剪、以獲得子網絡，并進行子網絡的微調訓練；從正則訓練、到子網絡微調，可協同知識蒸餾訓練，進一步保持子網絡的預測精度。主要步驟描述如下：
第一步：基于 Neuron與Attention head 的重要性，執行參數重排（Weight Re-wiring）：
第二步：多階段微調訓練，實現寬度、與深度方向的自適應正則化：
4.3 NAS-Bert
參考博客文章：
NAS-Bert——確保One-shot與Task-agnostic：https://blog.csdn.net/nature553863/article/details/120764285
4.4 Evolved Transformer
論文標題：
Lite Transformer with Long-Short Range Attention
收錄會議：
ICLR 2020
論文鏈接：
https://arxiv.org/abs/2004.11886
代碼鏈接：
https://github.com/mit-han-lab/lite-transformer
針對序列生成任務，基于 NAS 搜索獲得的 Transformer 結構:
搜索空間：包括兩個 Stackable cell，分別包含在 Transformer encoder 與 Transformer decoder。每個 Cell 由 NAS-style block 組成, 可通過左右兩個? Block 轉換輸入 Embedding、再聚合獲得新的 Embedding，進一步輸入到 Self-Attention 模塊。
搜索策略：基于 EA（Evolutional Aligorithm）的搜索策略；
網絡結構如下，融合了一維卷積與 Attention 的特點：
4.5 HAT: Hardware-aware Transformer
論文標題：
HAT: Hardware-Aware Transformers for Efficient Natural Language Processing
收錄會議：
ACL 2020
論文鏈接：
https://arxiv.org/abs/2005.14187
代碼鏈接：
http://github.com/mit-han-lab/hardware-aware-transformers.git
對 Transformer 模型進行網絡結構搜索時，通常會面臨兩個關鍵問題：
FLOPS 不能完全代表 Transformer 模型的計算速度/實際執行效率，即只能間接反映硬件平臺特性；
不同硬件設備上，與計算硬件或計算庫相適配的最優 Transformer 結構不盡相同;
針對序列生成任務，HAT（Hardware-aware Transformer）從搜索空間、搜索策略與搜索預測三方面加以分析：
搜索空間：Encoder-decoder Attention 的任意連接方式，以及 Transformer layer 內部結構（網絡深度、Attention head 數目、Width 等）；
搜索策略：訓練了 SuperTransformer 作為超網絡，并在超網絡預訓練之后，結合資源約束，通過進化搜索算法（EA: Evolutional Algorithm）尋找最優子網絡；不同于權重共享型 NAS（如 FBNet、SPOS 等），HAT 是一種動態可伸縮類型的One-shot NAS；
搜索預測：單獨、離線訓練了 MLP 回歸模型作為 Cost model 或 Predicter，用于預測不同子網絡結構（Architecture embedding）、在指定硬件平臺上的執行速度 （GPU 或 CPU latency），并且回歸預測效果良好（以相關性系數作為Metric）。搜索預測的好處在于：一方面可直接、有效獲取硬件平臺特性，作為超網絡預訓練的資源約束；另一方面，在訓練或搜索階段，無需在硬件平臺上測量推理延遲，提升搜索效率，從而有利于跨平臺模型結構搜索；
4.6 AutoFormer
論文標題：
AutoFormer: Searching Transformers for Visual Recognition
收錄會議：
ICCV 2021
論文鏈接：
https://arxiv.org/abs/2107.00651
代碼鏈接：
https://github.com/microsoft/AutoML
針對 Vision Transformer 的結構搜索，AutoFormer 提出的 Weight entanglement，在不額外增加 Block choice 的前提下，通過通道寬度、網絡深度、Attention head 數目等多個維度的調整，實現了 Vision Transformer 模型的動態可伸縮預訓練與結構搜索。
Weight entanglement 的做法，類似于 BigNAS、FBNet-v2 的通道搜索，都不會額外增加通道維度的權重參數量。相比于手工設計的 CNN 模型（ResNet、ResNext、DenseNet）與 Vision Transformer 模型（ViT、Deit），AutoFormer 模型在相同資源開銷條件下、能夠獲得最好的識別精度。
AutoFormer?總體思路如上圖所示，是一種基于 Weight entanglement 的動態可伸縮搜索方法，其搜索維度包括 Attention heads、通道寬度與網絡深度。下面從搜索空間、搜索策略與搜索效率這些維度加以分析：
搜索空間：Embedding dimension、Q-K-V dimension、Attention heads、MLP ratio 與 Network depth；根據不同的資源約束，分別設置 Supernet-tiny、Supernet-small 與 Supernet-base 三個基本的 Template model；
搜索策略：基于 Weight entanglement 的 One-shot NAS，完成超網絡預訓練之后，通過進化算法執行子網絡搜索。包含 l 個網絡層的子網絡結構可按下式表示，表示第 i 層的 Block 結構，表示相應的權重參數：
而每個 Block 結構，都是按照 Weight entanglement 原則從超網絡采樣獲得，下式表示 n 個動態選擇范圍：
搜索效率：按照權重共享形式、完成超網絡預訓練，訓練效率高、且收斂速度快；另外，由于沒有引入額外的 Block choice，因此訓練時的 Memory cost 較低。
計算圖優化
5.1 LINM
LINM（Loop-invariant Node Motion）是一種計算圖等價變換技術，通過將 Transformer 模型涉及的自回歸解碼（ Auto-regressive Decoding）的重復計算邏輯（Encoder-decoder Attention 的 k/v 計算、cache gather 等）移至 While-loop 之外，在確保計算功能不變的情況下，實現計算效率的提升：
5.2 MatMul替換為卷積
如果推理引擎的卷積算子具備很高的計算效率，并且為了避免一些不必要的訪存算子開銷，可以按如下計算圖等價變換，將 Dense layer 替換為 Conv1D layer：
5.3 QKV計算合并
MHA（Multi-head Attention）模塊在計算 Attention map 與 Context 編碼特征之前，需要基于輸入 Tensor 計算 QKV 矩陣、并 Split 成多份 Heads。可通過計算圖等價變換，將 QKV 計算合并到同一個 Dense layer，能夠充分利用 NPU 或 GPU 的并發計算特性：
1B,?N,?C?=?x.shape 2qkv?=?self.qkv(x).reshape(B,?N,?3,?self.sample_num_heads,?-1).permute(2,?0,?3,?1,?4) 3q,?k,?v?=?qkv[0],?qkv[1],?qkv[2]推理優化引擎
6.1 Faster Transformer
GitHub鏈接：
https://github.com/NVIDIA/FasterTransformer
Faster Transformer?是 NVIDIA 針對標準 BERT/Transformer 做的推理優化庫，其發布時間線：
● 1.0版本：
2019 年 7 月，開源了 FasterTransformer 1.0，針對? BERT 中的 Transformer Encoder 進行優化和加速；
面向 BERT 分類、自然語言理解場景；
底層由 CUDA 和 cuBLAS 實現，支持 FP16 和 FP32 計算，FP16 可充分利用 Volta 和 Turing 架構的 Tensor Core 計算單元；
提供 C++ API、TF Op 與 TensorRT Plugin 三種接口；
參考資料：https://zhuanlan.zhihu.com/p/73715272
●?2.0版本：
2020 年 2 月，新增對 Transformer decoder 的優化和加速，包括 decoder與 decoding 兩種加速模式；
面向生成式場景，如 NMT、文本內容生成與 ASR 等；
底層由 CUDA 和 cuBLAS 實現，支持 FP16 和 FP32 計算模，FP16 可充分利用 Volta 和 Turing 架構的 Tensor Core 計算單元；
提供 C++ API、TF Op 與 TensorRT Plugin 三種接口；
參考資料：快無止盡！FasterTransformer 2.0讓Decoder不再是你的性能瓶頸
●?2.1版本：
2020 年 6 月，引入 Effective Transformer 優化；通過 remove_padding 的支持，提高計算與訪存效率；
并新增 PyTorch Op 接口；
參考資料：https://github.com/bytedance/effective_transformer
●?3.0版本：
2020 年 9 月，新增 BERT encoder 的 INT8 量化加速支持；
僅支持 Turing 架構 GPU；
同時支持 PTQ 與 QAT 方法，提供了 TF 量化工具；
相比于 FP16 計算，約 20~30% 加速，但存在精度損失風險。
●?3.1版本：
2020 年 12 月，新增對 PyTorch 使用 INT8 推理的支援；
在 Turing 以后的 GPU 上，FP16 的性能比 3.0 提升了 10% ~ 20%；
INT8 的性能比 3.0 最多提升了 70%。
●?4.0版本：
2021 年 4 月，新增對 GPT-3 等百億/千億級參數規模模型的多機多卡推理加速能力；
新增 FP16 fused MHA 算子，同時支持 Volta 與 Turing 架構的 GPU；
以及對解碼端 Kernel 的優化，可以省略已完成語句的計算，節省計算資源；
參考資料：https://zhuanlan.zhihu.com/p/363517823；
●?支持矩陣：
Transformer 計算圖表示，主要包含三個部分的表示：
Encoder layer；
Decoder layer；
Decoding logic；
6.2 Torch-TensorRT
參考博客文章：
深度學習模型壓縮與優化加速（Model Compression and Acceleration Overview）
https://blog.csdn.net/nature553863/article/details/81083955
6.3 AICompiler
基于 AICompiler（AI 編譯優化），可通過 Op fusion 方式（算子融合），提升 Transformer 模型的執行效率。Op fusion 具備的好處主要是：1）首先，將多個零碎算子合并成計算功能等價的一個大算子，可減少密集的 Kernel launch 開銷 （CUDA 核函數啟動）；2）其次，由于相鄰算子之間存在內存讀寫開銷，Op fusion 可有效減少密集的內存訪問開銷；3）最終通過算子實現的自動尋優、與代碼生成（CodeGen），可自動生成優化的可執行代碼。?
總結
深度學習領域的高性能計算，總結為如下公式：
高性能計算 = 高效率算法 + 模型壓縮 + 系統/硬件優化
可以單獨的強調其中一種優化策略，或者實現多種選項的聯合優化（例如 Hardware-aware NAS，壓縮與編譯聯合優化等），以滿足實際應用場景的部署需求。在工業界，包括阿里淘系 MNN、阿里云 PAI、華為昇騰計算、Open AI Lab 的 Tengine、百度 EasyDL 等平臺，在該領域均構建了相對完善的軟硬件技術棧，為開放、普惠的 AI 應用提供了堅實基礎。
特別鳴謝
感謝 TCCI 天橋腦科學研究院對于 PaperWeekly 的支持。TCCI 關注大腦探知、大腦功能和大腦健康。
更多閱讀
#投 稿?通 道#
?讓你的文字被更多人看到?
如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。
總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。
📝?稿件基本要求：
? 文章確系個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標注?
? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發送，要求圖片清晰，無版權問題
? PaperWeekly 尊重原作者署名權，并將為每篇被采納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算
📬?投稿通道：
? 投稿郵箱：hr@paperweekly.site?
? 來稿請備注即時聯系方式（微信），以便我們在稿件選用的第一時間聯系作者
? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿
△長按添加PaperWeekly小編
🔍
現在，在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關注」訂閱我們的專欄吧
·
總結
以上是生活随笔為你收集整理的浅谈BERT/Transformer模型的压缩与优化加速的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： U盘显示用了很多但实际文件很小怎么回事
- 下一篇： 新年巨献！祝所有朋友新一年闪闪发光