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

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

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

論文題目：
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的基礎(chǔ)上進行優(yōu)化的。

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

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

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

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

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

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

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

初級配方：算法優(yōu)化

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

出處：

DeeBERT： Dynamic Early Exiting for Accelerating BERT Inference

（ACL‘20）

鏈接：

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

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

ACL'20的DeeBERT恰好就提出了一種動態(tài)的提前退出機制（Early Exit）。這一機制的設(shè)計是希望讓簡單的文本經(jīng)過較少的運算，而復(fù)雜的文本經(jīng)過更多的運算。

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

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

2. 動態(tài)注意力范圍

出處：

Adaptive Attention Span in Transformers

（ACL‘19）

鏈接：

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

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

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

具體地，mask函數(shù)基于兩個token之間的距離為注意力機制的權(quán)重計算添加了soft masking。注意力機制中的權(quán)重變?yōu)椋?/p>

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

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

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

3. 一階網(wǎng)絡(luò)剪枝

出處：

Movement Pruning： Adaptive Sparsity by Fine-Tuning

（NeurIPS‘20）

鏈接：

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

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

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

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

前向傳播過程中，神經(jīng)網(wǎng)絡(luò)利用加mask的參數(shù)計算輸出的各分量：。

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

對模型參數(shù)，有：

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

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

4. 零階網(wǎng)絡(luò)剪枝

出處：

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

（ICLR‘16）

鏈接：

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

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

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

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

5.動態(tài)浮點數(shù)量化

出處：

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

（arXiv Preprint）

鏈接：

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

誒，怎么還有？網(wǎng)絡(luò)的運算和存儲過程從深度、寬度和長度進行了全方位的優(yōu)化，模型還能再輕？

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

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

這正是AdaptivFloat的初衷：設(shè)計一種更加適合深度學(xué)習(xí)場景的浮點數(shù)數(shù)據(jù)類型。不過，要講明白AdaptivFloat數(shù)據(jù)類型需要涉及到一些與機器學(xué)習(xí)無關(guān)的知識。

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

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

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

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

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

高級配方：存儲介質(zhì)選擇

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

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

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

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

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

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

合并結(jié)果

好了，所有的基礎(chǔ)配方一個一個單獨使用的結(jié)果已經(jīng)出來了！那么，把它們?nèi)技釉谝黄鹉墚a(chǎn)生什么樣的結(jié)果呢？

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

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

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

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

究極配方：硬件加速器

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

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

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

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

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

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

總結(jié)

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

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

責(zé)任編輯：xj

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

文章出處：【微信公眾號：深度學(xué)習(xí)自然語言處理】歡迎添加關(guān)注！文章轉(zhuǎn)載請注明出處。