這個世界上有兩種極具難度的工程：第一種是把很平常的東西做到最大，例如把語言模型擴大成能夠寫詩寫文寫代碼的GPT-3；而另一種恰恰相反，是把很平常的東西做到最小。對于NLPer來說，這種“小工程”最迫在眉睫的施展對象非BERT莫屬。

從18年那個109M參數的BERT，到52M參數的蒸餾后的DistilBERT，再到14.5M參數的蒸餾更多層的TinyBERT，最后到12M參數的層級共享的ALBERT，曾經那個在集群上加載參數都費勁的BERT現在甚至已經可以跑在手機平臺上了。當我們為BERT的輕量化歡呼雀躍之時，有這樣一群人站了出來——只是手機端可不夠！他們的理想，是讓BERT跑在物聯網設備上，跑在低功耗芯片上，跑在我們能觸及的每一個電子器件上！

這樣一群來自哈佛/塔夫茨/HuggingFace/康奈爾的軟件和硬件極客們，此刻已披上了法袍，化身為為BERT極限瘦身的煉金術士，向著這個看似不可能的目標添加著許多讓人意想不到的配方…

論文題目：
EdgeBERT: Optimizing On-Chip Inference for Multi-Task NLP

論文鏈接：
https://arxiv.org/pdf/2011.14203.pdf

配方基底：ALBERT

出處：

ALBERT： A Lite BERT for Self-supervised Learning of Language Representations

（ICLR‘20）

鏈接：
https://arxiv.org/pdf/1909.11942.pdf

EdgeBERT是在ALBERT的基礎上進行優化的。

ICLR'20上谷歌提出的ALBERT是目前最佳的BERT壓縮方案。相比于過去利用知識蒸餾從原始BERT模型得到壓縮模型（例如DistilBERT [1]、TinyBERT [2]）和利用浮點數量化得到壓縮模型（例如Q8BERT [3]），ALBERT選擇直接拋棄BERT的預訓練參數，只繼承BERT的設計思想。正所謂不破不立，繼承BERT靈魂的ALBERT僅用12M參數就取得了和其他BERT變體相當的性能。

ALBERT對BERT的設計采取了以下三點改進：

嵌入層分解：BERT中，WordPiece的嵌入維度和網絡中隱藏層維度一致。作者提出，嵌入層編碼的是上下文無關信息，而隱藏層則在此基礎上增加了上下文信息，所以理應具有更高的維數；同時，若嵌入層和隱藏層維度一致，則在增大隱藏層維數時會同時大幅增加嵌入層參數量。ALBERT因此將嵌入層進行矩陣分解，引入一個額外的嵌入層。設WordPiece詞匯表規模為，嵌入層維度為，隱藏層維度為，則嵌入層參數量可由降低為。

參數共享：BERT中，每個Transformer層參數均不同。作者提出將Transformer層的所有參數進行層間共享，從而將參數量壓縮為僅有一層Transformer的量級。

上下句預測任務→句序預測任務：BERT中，除語言模型的MLM任務外還進行了上下句預測任務，判斷句2是否為句1的下一句，然而該任務被RoBERTa和XLNET等模型證實效果一般。作者提出將其替換為句序預測任務，判斷句2和句1之間的句子順序來學習文本一致性。

ALBERT的設計相當成功，成為了壓縮BERT的經典范例，而為了做到最極限壓縮BERT，以ALBERT作為起點確實是個不錯的主意。ALBERT已經如此強大，EdgeBERT又能壓到什么程度？作者一上來就用一張在QQP上的內存占用/運算量/性能比較圖吊足了讀者的胃口。（注意：縱坐標的內存占用量是對數刻度?。?/p>

本文對于ALBERT的利用除了作為初始化參數外，還在對下游任務做fine-tune時利用已經fine-tune好的ALBERT作為teacher進行知識蒸餾，來進一步提升模型性能。

初級配方：算法優化

1. 基于熵的提前退出機制

出處：

DeeBERT： Dynamic Early Exiting for Accelerating BERT Inference

（ACL‘20）

鏈接：

https://arxiv.org/pdf/2004.12993.pdf

ALBERT雖好，但Transformer太深了，算起來太慢，讓他變淺一點怎么樣？

ACL'20的DeeBERT恰好就提出了一種動態的提前退出機制（Early Exit）。這一機制的設計是希望讓簡單的文本經過較少的運算，而復雜的文本經過更多的運算。

實現上，DeeBERT向層的BERT模型添加了個“出口層”分類器（Early Exit Off-Ramps）。出口層分類器被放置在第和層Transformer之間，作為判斷第層Transformer的信息是否足以進行推斷的標志。進行推斷時，從最底層的分類器開始逐層計算出口層分類器的熵，當某一層的出口層分類器熵小于某個閾值時，則將該出口層分類器結果作為模型結果，省去后續層的計算。

上圖表現的是不同的熵閾值在MNLI、QQP、SST-2、QNLI四個數據集上的平均退出層數、理論運行時間節省和相應的準確度。添加提前退出機制后，在Acc損失1個百分點時，能夠在這四個數據集上分別降低30%、45%、54%和36%的理論運行時間；在Acc損失5個百分點時，能將在這四個數據集上的理論運行時間的降低進一步降低至44%、62%、78%和53%。

2. 動態注意力范圍

出處：

Adaptive Attention Span in Transformers

（ACL‘19）

鏈接：

https://arxiv.org/pdf/1905.07799.pdf

ALBERT雖好，但Attention范圍太廣了，算起來太慢，讓他變窄一點怎么樣？

ACL'19的Adaptive Attention提出的動態注意力范圍正是試圖通過這種方法減少注意力計算的。在Transformer的多頭自注意力機制中，不同head對應的注意力范圍并不一致，而讓每一個head都對所有token進行注意力運算無疑增添了無用開銷。為此，Adaptive Attention為每一個head添加一個不同的mask，使得每個token只能對周邊的token計算注意力，從而降低矩陣運算的開銷。

具體地，mask函數基于兩個token之間的距離為注意力機制的權重計算添加了soft masking。注意力機制中的權重變為：

其中為控制soft程度的超參數，為序列截止到token 的長度（原文采用了Transformer Decoder結構學習語言模型，故每個token只能于自己之前的token計算注意力。在EdgeBERT中沒有提及公式，不過根據模型圖的結構來看，分母應修改為對整個序列求和）。mask函數中的為mask的邊界，此邊界值會跟隨注意力的head相關參數和當前輸入序列變化：對于注意力機制中的每一個head ，有，其中、可訓練，為sigmoid函數。

EdgeBERT甚至對Adaptive Attention又進一步做了簡化：連算都不用算了，直接給每一個head賦一個可學習的，連輸入序列都不考慮了，多出來的參數只有12個（因為有12個head）。那么，這樣做的結果如何呢？作者將所有序列都pad/trunc到128長度，經過實驗，得到了一個驚人的結果：

表中展示的是經過優化后各個head的值，和模型在MNLI/QQP/SST-2/QNLI四個任務上的準確度。在一大半head幾乎完全被mask掉（）之后，模型居然只在這幾個任務上掉了0.5甚至0.05的準確度！而這一方法也為模型帶來了最高的計算量降低。

3. 一階網絡剪枝

出處：

Movement Pruning： Adaptive Sparsity by Fine-Tuning

（NeurIPS‘20）

鏈接：

https://arxiv.org/pdf/2005.07683.pdf

ALBERT雖好，但參數存起來占用的內存太長了，開銷太大，讓他變短一點怎么樣？

這里的網絡剪枝方式使用到了NeurIPS'20的一篇針對模型Fine-tune過程的剪枝算法。該論文的作者提出，傳統的零階網絡剪枝（即給模型里的參數絕對值設定一個閾值，高于它的保留，低于它的置零）的方法并不適用于遷移學習場景，因為該場景下模型參數主要受原始模型影響，卻需要在目標任務上進行fine-tune和測試，所以直接根據模型參數本身剪枝可能會損失源任務或目標任務的知識。與此相反，作者提出一種基于Fine-tune過程中的一階導數進行剪枝的Movement Pruning：盡可能保留fine-tune過程中更加偏離0的參數。

具體地：對于模型參數，為其賦予同樣size的重要性分數，則剪枝mask 。

前向傳播過程中，神經網絡利用加mask的參數計算輸出的各分量：。

反向傳播過程中，利用Straight-Through Estimator[4]的思想，將省略近似得到損失函數對重要性分數的梯度：

對模型參數，有：

將上述兩個式子代換后，省略的mask矩陣后可得：

根據梯度下降，當時，重要性增大，此時與異號。這表示，只有當在反向傳播時為正的參數變得更大或為負的參數變得更小時才會得到更大的重要性分數，避免被剪枝。

4. 零階網絡剪枝

出處：

Deep Compression： Compressing Deep Neural Networks with Pruning， Trained Quantization and Huffman Coding

（ICLR‘16）

鏈接：

https://arxiv.org/pdf/1510.00149.pdf

變短是變短了，但感覺這剪得還不夠好啊，再換一種算法讓它更短一點怎么樣？

這種方法的做法非常簡單：給模型里的參數設定一個絕對值閾值，絕對值高于它的保留，絕對值低于它的置零。由于方法實在太過簡單，不用公式也能很容易理解的吧(=?ω?=)

一階和零階網絡剪枝的效果對比如上圖所示（MvP：一階網絡剪枝，MaP：零階網絡剪枝）。在參數稀疏程度更高時，一階剪枝的效果更好，其他情況下是簡單的零階剪枝更有效。同時，研究還發現，當剪掉95%的嵌入層參數時，模型竟然在4個任務上都保持了至少95%的準確度。

5.動態浮點數量化

出處：

AdaptivFloat： A Floating-point based Data Type for Resilient Deep Learning Inference

（arXiv Preprint）

鏈接：

https://arxiv.org/pdf/1909.13271.pdf

誒，怎么還有？網絡的運算和存儲過程從深度、寬度和長度進行了全方位的優化，模型還能再輕？

的確，到此為止，一般的煉丹師已經看著自己三維裁剪過的模型感到成就感滿滿了，但這距離讓BERT跑遍所有設備的目標還差得遠。以下的部分就要開始深入一般NLP工程師見不到的第四個維度——硬件維度了。在深入硬件優化之前，先來道偏軟件的開胃菜，看看如何在浮點數量化上下手優化存儲吧！

我們平時在考慮利用浮點數特性進行計算加速時，最先考慮的是使用FP16混合精度，雖然有效，但畢竟損失了信息，效果也多多少少會受影響。既想保留精度又想加速訓練減小存儲，那就只有深入底層，修改浮點數的表示方法了！

這正是AdaptivFloat的初衷：設計一種更加適合深度學習場景的浮點數數據類型。不過，要講明白AdaptivFloat數據類型需要涉及到一些與機器學習無關的知識。

根據IEEE 754二進制浮點數標準，一個浮點數的二進制表示由三個域組成：符號位（Sign，）、指數偏移值（Exponent bias，）和分數值（Fraction，或Mantissa，）。由此，一個數可以表示為。

這時聰明的你可能發現有什么不對勁：按照表示無符號整形的方法取指數偏移值只能取出正數??！2的負次冪怎么辦！這正是為什么稱其為“指數偏移值”：它并不代表實際上2的指數，而是在其基礎上需要加一個常數作為2的指數：。

我們常用的浮點數保證了的選取能夠使得在數軸兩側分布幾乎均等（例如在32位浮點數FP32中，指數范圍為至），但這樣的數作為機器學習模型的參數顯然有些不太合適：為了增加小數的精度，我們甚至要允許這樣顯然不會出現的數也能表示，這真的不是在浪費內存？

AdaptivFloat的最關鍵動機正在于此：根據模型參數動態修改。所謂動態體現在每個Tensor都能得到量身定做的。方法也很簡單，找到Tensor里最大的一個數，讓它能被指數范圍恰好覆蓋到就好。不過說來簡單，為了實現這一方法需要配套地對現有浮點數表示方法進行許多其他修改，感興趣的話可以去看看AdaptivFloat原文，此外IEEE 754標準[5]同樣也可以作為參考哦~

上圖的結果中，Bit Width是浮點數的總位數，后五行在模型中使用了AdaptivFloat并將指數偏移值限制為3個bit。不比不知道，誰能想到修改個量化方式居然直接用8bit在四個數據集上都干翻了FP32？！不僅節省了3/4的內存，還取得了更好的性能，或許這就是吹毛求疵挑戰極限的極客們的浪漫？

高級配方：存儲介質選擇

光軟件優化可不夠！如果軟件優化就夠了的話，還買什么SSD，換什么GPU（不）

EdgeBERT的目的，是盡可能降低邊緣計算場景中使用BERT的推理延遲和耗能。為了最大限度地降低推理延遲，需要為網絡中不同的組成部分根據其增刪改查的需求選取符合最大效能的存儲介質。

BERT類模型的一大特點，在于它們都是預訓練模型：這類模型并非開箱即用，而是需要在目標任務上fine-tune后才能使用。這使得這類模型天生地存在著兩類存儲需求：

嵌入層：保存了Embedding向量。EdgeBERT在進行下游任務fine-tune時一般不對嵌入層進行修改。這類參數相當于只讀參數，只對快速讀取有較高要求，同時希望能夠在掉電時依然保持原有數據來降低數據讀寫開銷，因此適用耗能低、讀取速度快的eNVM（Embedded Non-Volatile Memory，嵌入式非揮發性記憶體）。本文選取的是基于MLC的ReRAM，一種低功耗、高速度的RAM。

其他參數：這些參數需要在fine-tune時進行改變。此處使用的是SRAM（與計算機內存的DRAM不同，SRAM更貴但功耗更低、帶寬更高，常被用于制造cache或寄存器）

嵌入層用了ReRAM究竟能帶來多大的影響？上圖結果表明，僅僅是改變了嵌入層的硬件介質就能帶來約的推理延遲降低，以及低至約的能耗降低！這對于邊緣計算場景下簡直是質變了?。楹蜶eRAM只有讀，但DRAM那邊卻要算DRAM讀+SRAM讀/寫呢？因為此處的ReRAM是特殊設計的只讀結構，并且可以直接讀入處理器進行運算。與此相反，DRAM，即電腦里一般使用的內存，需要經過基于SRAM的處理器cache，所以讀寫開銷需要加上這部分讀寫開銷。）

合并結果

好了，所有的基礎配方一個一個單獨使用的結果已經出來了！那么，把它們全都加在一起能產生什么樣的結果呢？

這張圖展示了完全體EdgeBERT在四個數據集上的性能、運算量和內存占用。其中所有紅點的實驗配置參照上表（即TABLE IV）。

在性能（準確度）相比ALBERT下降1個百分點時，EdgeBERT能取得的內存降低和的推理速度；下降5個百分點時甚至能取得的推理速度。

Embedding經過裁剪后僅保留了40%，使得存儲進eNVM的嵌入層參數僅1.73MB。

QQP的Transformer參數被mask掉80%，MNLI、SST-2、QNLI的Transformer參數被mask掉60%后，性能可以僅下降1個百分點。

究極配方：硬件加速器

這是什么東西？來，給你展示一下谷歌給樹莓派定制的TPU加速器Coral：

EdgeBERT專屬的硬件加速器想來應該也是差不多的樣子。

這一部分完全不是王蘇的菜…給各位上一張EdgeBERT加速器的硬件結構圖：

感興趣的各位可以去參照原文進行學習_(:з」∠)_

這個加速器有什么用呢？它是基于EdgeBERT的運算特點量身定做的加速器，能夠把fine-tune好的EdgeBERT完整地裝進去進行運算。至于運算效果，它們修改了模型圖中VMAC序列（即進行矩陣運算的單元序列）長度，與NVIDIA的移動端TX2 mGPU進行了推理時間和耗能的對比：

本文中提出的硬件加速器能夠為EdgeBERT帶來相比于baseline硬件加速器的能耗降低，相比于英偉達TX2移動端GPU甚至能夠帶來的能耗降低！耗電大戶BERT家族終于也有能被說“省電”的一天了！

總結

壓縮BERT是一項研究，但極限壓縮BERT則是一項不易完成的工程：無論是對Transformer模型的全方位裁剪，還是對硬件存儲介質讀寫性能與容錯的取舍，亦或是對專屬硬件加速器的設計，單獨拿出一項已足夠艱難，將它們合在一起就不僅可能互相沖突，甚至還可能產生逆向優化。這篇文章通過大量的實驗，測試了已有的幾種優化方法在邊緣計算場景下的性能，比較了不同優化方法之間的差別，分析了所有優化方法進行組合后的影響及效果，并進一步提出了專屬的硬件結構，實現了對目前已有的最輕量BERT變體的即插即用。對于需要長待機、低功耗、短延遲的場景，例如智能家居或是其他需要NLP技術加持的物聯網設備，或許我們真的能在不遠的將來看到實體的類似EdgeBERT加速器的解決方案出現。

雖然我們對于可能帶來更大變革的模型結構依然處在探索當中，但從當下實用的角度而言，用基于Lottery Ticket Hypothesis[6]的BERT優化方法尋找一個更優的類BERT子結構依然是一個不錯的課題，至少它能讓更多人、更多時候、更多場景能夠用上性能強大的預訓練模型。本文中提到的這些優化方法是不是也給愛思考的你帶來了什么啟發呢？

責任編輯：xj

原文標題：軟硬兼施極限輕量BERT！能比ALBERT再輕13倍？！

文章出處：【微信公眾號：深度學習自然語言處理】歡迎添加關注！文章轉載請注明出處。