作者 |
姚益武
單位 |
阿里巴巴集團
研究方向 |
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
參考部落格文章：
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 任務，
AdaBer
t
透過可微分搜尋、與知識蒸餾，實現卷積型別的 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
參考部落格文章：
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 的併發計算特性：
推理最佳化引擎
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）
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」
·