基于动态提前退出的预训练模型推理加速简介
?PaperWeekly 原創 ·?作者?|?李磊
單位?|?北京大學碩士生
研究方向?|?自然語言處理
隨著如 BERT、RoBERTa 等預訓練模型在一眾自然語言處理任務上取得了出色的表現,如何在取得較好性能的前提下,縮小模型的大小,加快模型的推理速度,成為了學術界和工業界的一個熱門的研究方向。這篇文章將會介紹基于動態提前退出( Dynamic Early Exiting) 的預訓練模型推理加速方面的系列工作。
動態提前退出
Dynamic Early Exiting(動態提前退出)的核心是根據樣本的難度來匹配相應的計算量,這一思想在之前的不少文獻中都已有體現,例如在 RNN 中根據樣本難度決定所需的時間步(ACT),CV 領域的 Shallow-Deep Network/ 動態深度網絡等。將這一思想應用到如 BERT 的預訓練的模型上,我們可以把模型的 Layer 視作是一個固定的時間步/深度,則 BERT(base)實際上對每個樣本都同等地執行了 12 層的前向計算。
然而,樣本之間是普遍存在難度差異的,對于一些比較簡單的樣本,或許前幾層的 BERT Layer 表示就足以支撐模型進行正確的預測,則后續的計算都可以認為是冗余的,從而可以省去以提升推理的速度。
基于這個思想,一個最為直接想法躍然紙上,我們可以在每個層之后增加一個線性分類器,然后每執行一層模型,就拿著得到的表示去過分類器得到輸出結果,并且根據這個輸出結果來判斷是否可以直接退出,或者繼續執行下去,直到模型的最后一層,整個框架的示意圖如下:
框架有了,那么如何進行是否要退出模型的決策呢?發表在 ACL 2020 上的 DeeBERT [3] 采用了一個非常啟發式的指標,即模型預測概率分布的熵,若其熵小于某個特定的閾值,說明模型對于當前的預測足夠自信,則可以 退出,否則則需要繼續執行下一層,整體的算法如下:
最后則就是中間的這些分類器應該如何訓練,DeeBERT 采用的是兩階段的方案,第一個階段就是簡單的 Fine-tuning,訓練模型每一層的表示以及最后一層的分類器;第二階段則固定住中間層的參數以及最后分類器,利用 CE loss 訓練中間層的分類器,從而優化中間分類器的分類結果。
模型的評估則是通過設定不同的閾值,比較不同加速比之下的測試集上的平均性能。作者在 GLUE 的分類數據集上進行了實驗,發現相比于 DistilBERT 這類靜態模型壓縮的方法,能夠取得更好的推理加速的效果:
無獨有偶,同樣是 ACL 2020,還有兩篇也是基于動態退出的想法進行 BERT 推理加速的文章,這里對其簡要介紹如下:
FastBERT [5] :和 DeeBERT 不同之處在于其中間層的分類器用了更加復雜的基于 attention 的分類器,并且在訓練中間分類器的時候增加了額外的 self-attention 目標,使得中間層的輸出和最后一層分類器的輸出保持一致。文章在一種文本分類數據集能夠在 ~10x 加速比的情況下保持原始模型的性能。
RightTool [4] :文章利用中間分類器的預測結果的 confidence,即預測類別的概率值作為退出的指標,并且對其進行了 post-hoc carlibration,以避免 over-confident 的問題。文章和之前的 Paper 不同的地方在于,其對于樣本復雜度進行了分析,發現樣本長度和難度的相關性很低,而模型的預測的置信度和難度的相關性也不高,相關系數大約在 0.3 附近,這也提示我們說雖然目前基于啟發式的退出策略已經取得了不錯的成績,但還有著提升的空間。
改進
后續有不少工作嘗試對 Dynamic Early Exiting 的框架進行探究和改進,CascadeBERT [6] 總結了 Dynamic Early Exiting 在高加速比下,即大部分樣本都在淺層退出時,會遇到的兩個瓶頸問題:(1)退出決策不魯棒,即樣本難度和中間分類器的分類置信度并不匹配;(2)淺層表示具備的信息不夠充分。
對于前者,作者設計了一個用于衡量退出決策和樣本難度順序是否匹配的指標 Difficulty Inversion Score(難度逆序分數),越高說明中間層分類器給出的退出決策和樣本難度越一致,文章發現 DeeBERT 的中間分類器在其自己 選擇退出的樣本上進行了分析,發現其依舊無法給出較好的退出決策:
對于后者,作者比較了使用相同層的表示,并且在 DeeBERT 選擇退出的樣本上進行了性能分析:
發現即使是相較于直接 fine-tune BERT 前幾層得到的 BERT-kl,DeeBERT 的性能在只用少數淺層的樣本的時候也難以做出相對準確的預測。基于此,CascadeBERT 提出了一種級聯不同大小預訓練模型的方案,并且在訓練過程中加入樣本難度相關的正則項,從而使得在較高加速比的情況下也能夠保持較好的性能。
同期也有不少工作在這兩方面對 Dynamic Early Exiting 進行改進:
2.1 更好的退出決策
PABEE [7] 嘗試結合中間不同層的退出決策的一致性進行更加魯邦的退出決策,其核心是當連續某幾層的中間分類器都給出一致的預測結果后,認為這一個樣本的預測結果以及足夠自信,因而可以提前退出:
這樣的方案簡單并且有效,相對于 DeeBERT 都取得了比較明顯提升,但因為在設定退出條件的時候需要設置的一致預測層數只能是整數,大大地限制了其加速比的取值范圍。類似地,ELBERT [8] 在連續多層分類器給出某個類別上單調遞增的概率之后觸發退出。
Early Exiting with Ensemble [9] 指出目前的動態退出框架距離性能的上限仍然存在巨大的空間,假設存在 oracle 模型能夠在最早的完成正確預測的一層退出,這一情況即下的加速比即是理論的一個上限,而目前不同的方法還存在很大的改進空間,下圖的右上角是理論的上界,而大部分的方法都落在了圖中的左下部位:
作者基于 Ensemble 理論,推導出一個目標函數,使得不同層的 classifier 在正確的類別上預測一致,而錯誤類別上的分布盡量拉大,進而增加了中間分類器預測的 Diversity,使得后續基于 Ensemble Voting 的退出決策更加準確。在多個數據集上的實驗結果也證明了這一點:
類似地,LeeBERT [10] 使中間的分類器進行 Mutual Distillation,并且對不同層的分類器的 loss 賦以不同的可學習的權重,以提升退出決策的質量。
2.2 更好的中間表示
前面提到,CascadeBERT 指出 BERT 淺層的樣本表示可能不足以支撐分類器進行正確的分類決策,因其,其作者提出將中間層退出改為在不同大小的預訓練模型(例如,一個 2 層的 BERT 和一個 12 層的 BERT)之中進行級聯退出,因為 2 層的 BERT 也是從頭預訓練得到,因而其樣本表示具備比較完整的語義信息,從而能得到一個比較好的效果:
Global Past-Future?[11] 則嘗試利用 imitation learning,一方面利用所有淺層的樣本表示,另外一方面嘗試預測出更深層的樣本表示來作為輔助信息,進而提升分類的效果:
應用
動態提前退出不僅僅可以應用在簡單的分類任務上,不少研究也嘗試將其拓展到更加不同的任務上。
BERxiT [12] 利用一個額外學習 learn-to-exit 模塊,根據當前的樣本表示給出一個是否要退出的決策,這一就可以將這一思想應用到回歸任務上。
Early Exiting for Sequence Labeling [13] 更近一步地拓展到了序列標注任務上,其根據序列標注任務的特點,設計了兩種退出的機制:Sentence-level Early Exit (SENTEE)以及基于局部 token 的不確定性設計的 Token-level Early Exit(TOKEE):
句子級別的退出機制和 DeeBERT 類似,就不再贅述。Token 級別的退出,則是基于上下文窗口 中所有 Token 的預測不確定度的最大值,若是其超過了某個設定的閾值,則后續的計算過程中,這一 token 的表示將不再被重復計算,而是直接傳遞到下一層,進而能夠將 A ttention計算量從 到 , 和 分別是序列長度和沒有退出的 token 的數量, 是隱層的維度。
SENTEE 和 TOKEE 能夠在一眾序列標注任務上取得比較高的加速比的同時保持較好的效果:
小結
我們在這篇文章中簡要介紹了用于預訓練模型推理加速的動態提前退出思想,并且對目前主要的工作和相關改進進行了梳理。此外,這一框架也被拓展到例如序列標注等的多種任務上,展現出了一定的潛力,也期待未來有更多相關工作來探索更好更快的預訓練模型。
更多相關的 paper 可以參考復旦大學博士生孫天祥同學整理的 Paper List:
https://github.com/txsun1997/awesome-early-exiting
我們也梳理了和預訓練模型效率相關的 Paper List:
https://github.com/TobiasLee/Awesome-Efficient-PLM
包括剪枝(Pruning)、知識蒸餾(Knowledge Distillation)、量化(Quantization)等,也歡迎大家 Star。
//預訓練模型專題交流群//
隨著算力的提升,大數據的涌現,新一代人工智能已經邁向“煉大模型”的時代。如果你也對預訓練模型感興趣,歡迎掃描下方二維碼(備注「大模型-單位/學校-研究方向」)加入群聊,聽說這里還可以免費體驗 1060 上車百億大模型推理喲~
掃描下方二維碼
參考文獻
[1] Adaptive Computation Time for Recurrent Neural Networks. Preprint 2016
[2] Shallow-Deep Networks: Understanding and Mitigating Network Overthinking. ICML 2019
[3] DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference. ACL 2020.
[4] The Right Tool for the Job: Matching Model and Instance Complexities. ACL 2020.
[5] FastBERT: a Self-distilling BERT with Adaptive Inference Time. ACL 2020.
[6] CascadeBERT: Accelerating Inference of Pre-trained Language Models via Calibrated Complete Models Cascade. Findings of EMNLP 2021.
[7] BERT Loses Patience: Fast and Robust Inference with Early Exit. NeurIPS 2020.
[8] ELBERT: Fast Albert with Confidence-Window Based Early Exit. Preprint 2021.
[9] Early Exiting with Ensemble Internal Classifiers. Preprint May 2021.
[10] LeeBERT: Learned Early Exit for BERT with Cross-Level Optimization. ACL 2021.
[11] A Global Past-Future Early Exit Method for Accelerating Inference of Pre-trained Language Models. NAACL 2021.
[12] BERxiT: Early Exiting for BERT with Better Fine-Tuning and Extension to Regression. EACL 2021.
[13] Accelerating BERT Inference for Sequence Labeling via Early-Exit. ACL 2021
特別鳴謝
感謝 TCCI 天橋腦科學研究院對于 PaperWeekly 的支持。TCCI 關注大腦探知、大腦功能和大腦健康。
更多閱讀
#投 稿?通 道#
?讓你的文字被更多人看到?
如何才能讓更多的優質內容以更短路徑到達讀者群體,縮短讀者尋找優質內容的成本呢?答案就是:你不認識的人。
總有一些你不認識的人,知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁,促使不同背景、不同方向的學者和學術靈感相互碰撞,迸發出更多的可能性。?
PaperWeekly 鼓勵高校實驗室或個人,在我們的平臺上分享各類優質內容,可以是最新論文解讀,也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個,讓知識真正流動起來。
📝?稿件基本要求:
? 文章確系個人原創作品,未曾在公開渠道發表,如為其他平臺已發表或待發表的文章,請明確標注?
? 稿件建議以?markdown?格式撰寫,文中配圖以附件形式發送,要求圖片清晰,無版權問題
? PaperWeekly 尊重原作者署名權,并將為每篇被采納的原創首發稿件,提供業內具有競爭力稿酬,具體依據文章閱讀量和文章質量階梯制結算
📬?投稿通道:
? 投稿郵箱:hr@paperweekly.site?
? 來稿請備注即時聯系方式(微信),以便我們在稿件選用的第一時間聯系作者
? 您也可以直接添加小編微信(pwbot02)快速投稿,備注:姓名-投稿
△長按添加PaperWeekly小編
🔍
現在,在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關注」訂閱我們的專欄吧
·
總結
以上是生活随笔為你收集整理的基于动态提前退出的预训练模型推理加速简介的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 肥羊爱上大嘴鱼怎么做?
- 下一篇: 佛山大沥步行街到黄飞鸿街怎么坐地铁?