降低預測過程計算成本,這些NLP模型壓縮方法要知道

降低預測過程計算成本,這些NLP模型壓縮方法要知道

編譯 | 凱隱

出品 | AI科技大本營(ID:rgznai100)

近年來,基於谷歌Transformer的語言模型在神經機器翻譯,自然語言推理和其他自然語言理解任務上取得了長足進展。

通過多種語言模型的平均損失進行自我監督預訓練,使得在大範圍的語料庫上訓練的模型可以在許多任務中提高下游性能。然而,大量的參數和計算量仍然是阻礙BERT和其衍生模型部署的難點。

值得慶幸的是,在過去的兩年裏,我們已經看到了各種各樣的技術,可以縮短模型在進行實際預測時消耗的時間。因此,這篇文章主要着眼於在基本模型預訓練後可以用於降低預測過程計算成本的方法,主要包含以下方法:

  • 數值精度約簡:通過降低計算過程中使用的浮點數精度(浮點約簡)和量化,來加速計算。
  • 計算融合:在計算圖中選擇節點並進行合併的技巧。
  • 網絡修剪: 識別和刪除網絡中不重要的部分。
  • 知識提煉: 訓練更小的,效率更高的模型來模仿表現力更強,同時計算成本更高的大模型。
  • 模塊替換:通過替換部分模塊來降低模型的深度和複雜度。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

數值精度約簡

數值精度約簡可能是爲模型帶來預測加速最通用的方法。在過去幾年GPU對16位浮點操作的支持性很差,這意味着降低權重和激活值的精度往往不會帶來加速,甚至造成減速。英偉達(Nvidia)Volta和圖靈張量核架構的引入在一定程度上解決了這個問題,讓GPU能夠更好的實現高效的16位浮點精度運算。

1、浮點數的表示

浮點類型數據主要存儲三種類型的數值信息:符號,指數,分數。傳統的32位浮點表示法分別用8位和23位來表示指數和分數(剩下一位應該表示正負),而傳統的16位表示法(用於NVIDIA硬件的格式)將32位表示法中的指數和分數部分大致減半。相對於GPU而言,TPU(張量處理單元)則使用一種稱爲bfloat16的變體,它可以選擇將一些位從分數移動到指數,即犧牲一定的精度來換取表示更大範圍的值的能力。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

三種不同的浮點數表示方法

Transformer網絡的大部分結構都可以直接轉換爲16位浮點數的權重和激活值,且不會帶來計算準確率的下降。而網絡的一小部分,特別是softmax操作部分必須保留爲32位浮點精度。這是因爲大量的小數值(我們的logits)累積起來就可能成爲錯誤的來源。由於這樣做同時使用了16位和32位精度的浮點數值,因此這種方法通常稱爲“混合精度”。

低精度的數值表示主要可以從兩方面實現加速:

(1)機器自帶的半精度指令(速度更快)

(2)更大的batch size(得益於更緊湊的表示)

NVIDIA已經發布了一套相當廣泛的與浮點精度降低相關的基準測試——實際上,這種方法可以使速度提高3倍。

鏈接:https://github.com/NVIDIA/DeepLearningExamples

2、整型量化

32位浮點到8位整數型值的量化也能實現加速,但需要更細節的實現。特別的,爲了確保8位整型的計算結果與32位浮點計算結果儘可能相同,還需要在訓練後對輸入進行校準。

如果已經知道網絡激活值可能所佔用的範圍,可以將該範圍劃分爲256個離散塊,並將每個塊分配給一個整數。只要存儲了比例因子和佔用的範圍,就可以使用整數近似來進行矩陣乘法,並將乘法輸出恢復爲浮點值。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

兩種不同的整型量化方法

一種簡單的整型量化方法,如上圖右側,首先選擇一個範圍和偏移量,這樣就不會將校準後的輸入浮點激活值激活映射到unit8值(-128,127)兩端的整數。但這種做法在適應兩端極值的同時犧牲了一些精度。如上圖左側,爲了解決這個問題,TensorRT之類的框架會選擇合適的scale和offset值,以最小化模型的32位浮點版本和8位整型版本的輸出激活值之間的KL差異。這使我們能夠以合理的方式在取值範圍和計算精度之間權衡取捨。由於KL散度可以被看作激活值在不同編碼(不同精度可以看做不同編碼)下的信息損失的度量,因此是一種自然合理的度量方式。

關於如何使用英偉達TensorRT來將8位整型量化應用到個人模型的更多詳細細節,可以參考一下鏈接:

https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html

https://github.com/tensorflow/tensorrt

3、網絡層融合與計算圖優化

除了浮點數約簡和整型量化外,操作融合提供了一個實用的,通用的選擇來實現更高效的預測。操作融合的基本原理是將計算流程中若干層所執行的操作進行合併,以避免對全局設備存儲器的冗餘訪問,進而提高執行效率。通過將多個計算操作合併到一個內核中,可以加速內存讀寫速度。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

操作合併示例

如上圖,我們可以將添加的skip連接與層規範化操作(LN)的比例(scale)和偏差(bias)合併在一起。

軟件優化允許我們重構一些矩陣乘法操作以更好地進行計算並行話。如下圖,我們可以將self-attention層的查詢向量(Query)、鍵向量(Key)和值向量(Value)的計算投影合併到一個矩陣乘法中。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

對自注意力層進行優化示意圖

遺憾的是,從這類計算流程圖優化中很少看到關於加速幅度的詳細信息,但這種改進仍然是有效的,大約有10%的吞吐量提升。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

網絡修剪

除了純粹的軟件層面優化,還有其他許多方法可用來修剪網絡結構,刪除對最終模型貢獻小的權重。許多剪枝方法(如Fan等人提出的"Reducing Transformer Depth on Command With Structured Dropout")需要在訓練前對網絡進行修改,以生成足夠稀疏的模型,並在訓練後進行剪枝。剪枝文獻中的其他論文聚焦於在沒有具體預測目標的情況下,模型所學習的連接模型如何才能更加稀疏。(例如Gorden等人提出的"Compressing BERT: Studying the Effects of Weight Pruning on Transfer Learning")

雖然這些方法本身都很有意義(而且結構化的刪除部分層的方法在實際應用中有重要前景),但更有趣的是一種能以特定方式應用並且仍然能獲得性能提升的方法。這類方法基於一個公認事實來對模型進行修剪,即解決特定任務只需要模型的一部分。

爲了獲得經驗上的模型性能提升而進行修剪需要結構化的稀疏性。簡單地將單個權值歸零並不足以獲得性能提升,因爲我們沒有實際的方法來利用這種稀疏性。因此,我們必須設法砍掉網絡中的大部分結構,以獲得實際的性能提升。

注意力頭剪枝

在文獻"Are Sixteen Heads Really Better than One?"中,作者通過迭代的方法從BERT模型中逐步去除注意力頭(attention head)。他們使用了一種基於梯度檢測的方法(對下游任務進行梯度估計)來估計每個注意力頭的重要性,並通過繪製性能--去除的注意力頭所佔百分比函數來測試模型對注意力頭剪枝的魯棒性。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

如下圖,在實踐中,作者發現20 - 40%的注意力頭可以修剪,且對模型準確性的影響可以忽略不計。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

基於門控的模型修剪

在J.S. McCarley,Rishav Chakravarti和Avirup Sil合著的《Structured Pruning of a BERT-based Question Answering Model》一書中,坐着探索了一種更通用的模型剪枝方法。不僅關注注意力頭,還對每一層的輸入以及每一BERT層的前饋層的激活值進行了門控。

他們探索了一些距離度量機制,從而有選擇的對網絡進行修剪,這包括Michel等人提出的重要措施。但最終確定了一個L0正則化項(添加到損失函數中),並且該項可以在超參調優時進行設置(作爲超參數之一),來提高模型的稀疏性。爲了使這個L0正則化項可微,他們使用了一個類似於在變分自動編碼器中使用的重參數化技巧。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

在他們的實驗中他們發現通過正則化項進行稀疏懲罰,並通過參數調優得到的結果優於“Are 16 Heads Really Better than 1”使用的基於重要性的估計方法,並且他們發現可以另外刪除近50%的前饋激活值,且在基準任務上對最終性能影響微乎其微。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

爲了獲得更多提升,作者還使用了知識蒸餾技術。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

知識蒸餾

1、發展歷史

知識蒸餾首先由Geoffrey Hinton, Oriol Vinyals, 和 Jeff Dean 在2015年的工作"Distilling the Knowledge in a Neural Network"中提出。知識蒸餾主要通過損失修正的方法,將教師模型中包含的知識遷移到學生模型中。

首先,假設我們可以訪問大量未標記的樣本。如果我們相信教師模型的預測結果是足夠準確的,但是教師模型在實際應用中部署使用太麻煩或計算成本太昂貴,我們可以使用教師模型來預測未標記樣本池中的目標類,並將這些目標作爲監督樣本用於學生模型的訓練。如果不是生成與最大可能性類相對應的硬目標,而是生成所有可能類的概率分佈,那麼學生模型將獲得更多的信息豐富的監督樣本。

直觀來講,學生模型的一些預測錯誤比其他錯誤更加合理,如下圖,將哈士奇誤認爲勺子是荒謬的,但是將哈士奇誤認爲雪橇犬則是合理的。即設置的損失函數應該反映錯誤的嚴重程度。通過懲罰教師模型預測結果和學生模型預測之間的差異(即鼓勵二者預測結果相同),學生網絡可以從教師網絡給出的預測中學習更多有意義的信息。在語音識別任務中,教師網絡的大部分網絡性能可以通過僅約3%的訓練數據來重現。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

有證據表明高參數量可能對樣本的有效學習至關重要,而且,就相同時間而言,以一個固定的複雜度訓練大型模型可能比訓練一個等價的緊湊模型效果更好。因此有效地把教師模型學到的知識傳遞給被壓縮的學生模型是很有前途的。

2、相似模型的知識遷移

在上面討論文章"Structured Pruning of a BERT-based Question Answering Model"中,通過知識蒸餾將未修剪的教師模型中包含的知識傳遞給修剪過的學生模型。在自然問題數據集上,教師模型在長答案和短答案問題上的F1分數分別爲70.3和58.8。隨着大約50%的注意力頭和前饋激活值的削減,F1分說分別下降到67.8和55.5,即大約有2.5的下降。但如果用蒸餾損失來代替超參調優過程中的交叉熵損失,那麼性能將恢復1.5到2個點,達到69.3和58.4。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

不同的模型壓縮方法帶來的性能損失對比

由Victor Sanh,Lysandre starter,Julien Chaumond和Thomas Wolf提出的“DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter”,在一個語言建模任務的二級預訓練步驟中,從基本BERT模型向一個6層的BERT壓縮後的學生模型執行知識蒸餾。學生模型(以任務不可知的方式訓練)在GLUE基準測試上保持了97%的模型性能,同時減少了60%的預測時間。

在文章“TinyBERT: Distilling BERT for Natural Language Understanding”中,作者採用了從BERT模型到一個4層,隱藏神經元個數爲312的學生模型的知識蒸餾與遷移。他們在預訓練和調參時都進行了遷移,得到的模型在GLUE基準測試上的達到了BERT-base性能的96%,且模型相對於BERT-base縮小了7.5倍,預測推斷速度提高了近10倍。

在文章"Patient Knowledge Distillation for BERT Model Compression"中,作者將知識蒸餾損失函數應用於12層BERT教師模型以及6層學生模型,這在大約5/6的GLUE任務中得到了精度提升(相對於僅將知識蒸餾應用於模型本身)。

3、不同結構模型的知識遷移

在目前討論的論文中,教師模型和學生模型具有相同的基本架構,學生模型通常使用教師模型的權重來進行初始化。然而,即使在教師和學生模型結構差異很大的情況下,也可以應用知識蒸餾損失來讓二者的預測結果相近,從而將教師模型學習到的知識遷移到學生模型。

在文章"Training Compact Models for Low Resource Entity Tagging using Pre-trained Language Models"中,作者首先在命名實體識別任務上訓練了一個BERT教師模型(參數量約330M),然後將其遷移到一個更緊湊更高效的CNN-LSTM學生模型上(參數量約3M),這樣做使得他們在CPU硬件上以最小的精度損失實現了高達2個數量級的速度提升。

在文章"Distilling Transformers into Simple Neural Networks with Unlabeled Transfer Data"中,作者將BERT-Base和BERT-Large遷移到一個雙向LSTM學生模型上,使得學生模型在4種分類任務(Ag News, IMDB, Elec, and DBPedia)上都能達到和教師模型相同的準確度,且參數量削減爲13M。他們還發現蒸餾使得樣本效率大大提高,每個任務只需要用大約500個帶標籤的樣本進行訓練,就能讓學生模型達到和教師模型相同的水平(提供足夠的未標記樣本進行測試)。

在文章"Distilling Task-Specific Knowledge from BERT into Simple Neural Networks"中,作者Lin報告了在使用參數小於1M的單層BiLSTM執行各種句子配對任務(QQP、MNLI等)時的類似結果。

在文章"Attentive Student Meets Multi-Task Teacher: Improved Knowledge Distillation for Pretrained Models"中,作者將多任務學習與知識蒸餾相結合,將Transformer教師模型通過注意力機制遷移到deep LSTM學生模型上。文章指出,從知識蒸餾中得到的提升與多任務學習框架帶來的泛化提升是一樣的,並且預測速度是純粹知識精餾的30倍,是TinyBERT的7倍。

知識蒸餾是最近很流行的方法,原因很明顯——它很可能成爲許多基於Transformer的語言模型應對逐漸增加的參數量的有效方法。如果我們想要儘可能地利用GPU,那麼我們就需要通過知識蒸餾這樣的方法來保持高預測速度。

模塊替換

將要介紹的最後一篇論文有兩個特點,一是採用了一種更新穎的模型壓縮方法,二是與下面展示的現代藝術作品一起發表:

降低預測過程計算成本,這些NLP模型壓縮方法要知道

“BERT-of-Theseus:Compressing BERT by Progressive Module Replacing”是由Canwen Xu等提出的工作,不同於之前的工作僅訓練一個單獨的學生模型來最小化知識蒸餾損失,BERT-of-Theseus在調試階段隨機地用一個新的模塊(後繼模塊)來替換原來的模塊(前繼模塊),且在每個訓練批次都隨機替換。

BERT-of-Theseus是思想實驗“Ship of Theseus”的一個版本,該實驗探究的是一艘船經過一點一點的修復和升級後,是否仍然是原來的物體。BERT-of-Theseus將這種逐步替換思想應用到模型壓縮的思想中。

後繼模塊通常是前繼模塊的廉價版本——在本例中,單個變壓器層替換了一個由2個變壓器層組成的塊。與知識蒸餾不同,這裏沒有使用損失來鼓勵後續模塊模仿它們的前輩,而只是通過簡單的讓繼任模塊和前任模塊可以互換使用,使得繼任者學習模仿前任的行爲。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

這種隱式模仿行爲的一個優點是,我們不再需要選擇如何重量各種知識蒸餾損失目標模型的損失,不同於知識蒸餾使用的L0正則化損失 ,不存在二次訓練的步驟,與下游整合壓縮併發執行。最後,漸進式模塊替換方法也適用於整個模型結構,且它在設計中沒有利用Transformer模型的任何特定特性,因此泛化能力更強。

作者用線性學習率變化策略進行了實驗,發現隨着時間的推移,線性增加的模塊替換率比恆定的替換率效果更好。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

爲了測試這種方法的魯棒性,作者在將BERT-base應用到GLUE基準測試時使用了Theseus壓縮,並且輕鬆地超過了其他幾個基於知識蒸餾的方法,這些方法在將原始模型壓縮到50%的大小時,常常落後於BERT-base不到1個點。

降低預測過程計算成本,這些NLP模型壓縮方法要知道

漸進式模塊替換很有吸引力的部分原因是它提供了一種新的訓練方法,可以用其他方法進行試驗來增加模型吞吐量,而這些方法通常需要對scratch進行重新訓練。對於獨立的研究人員和較小的公司來說,從頭開始重新培訓transformer模型通常是很困難的,因此很難利用那些提出了提高模型效率的有用想法但沒有發佈預先培訓過的模型的論文(即沒有開源的論文)。

論文地址:

https://arxiv.org/pdf/2002.11985.pdf

http://mitchgordon.me/machine/learning/2019/11/18/all-the-ways-to-compress-BERT.html

降低預測過程計算成本,這些NLP模型壓縮方法要知道

降低預測過程計算成本,這些NLP模型壓縮方法要知道

  • 那個分分鐘處理10億節點圖計算的Plato,現在怎麼樣了?

  • 看似毫不相干,哲學與機器學習竟有如此大的交集

  • 黑客用上機器學習你慌不慌?這 7 種竊取數據的新手段快來認識一下

  • “谷歌殺手”發明者,科學天才 Wolfram

  • 數據庫激盪40年,深入解析PostgreSQL、NewSQL演進歷程

  • 5分鐘!就能學會以太坊 JSON API 基礎知識

相關文章