EdgeBERT:极限压缩bert
這個世界上有兩種極具難度的工程:第一種是把很平常的東西做到最大,例如把語言模型擴大成能夠寫詩寫文寫代碼的 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 因此將嵌入層進行矩陣分解,引入一個額外的嵌入層 E。設 WordPiece 詞匯表規模為V,嵌入層維度為 E,隱藏層維度為 H,則嵌入層參數量可由 O(V x H)降低為O(V x E + E x H)。
參數共享:BERT 中,每個 Transformer 層參數均不同。作者提出將 Transformer 層的所有參數進行層間共享,從而將參數量壓縮為僅有一層 Transformer 的量級。
上下句預測任務→句序預測任務:BERT 中,除語言模型的 MLM 任務外還進行了上下句預測任務,判斷句 2 是否為句 1 的下一句,然而該任務被 RoBERTa 和 XLNET 等模型證實效果一般。作者提出將其替換為句序預測任務,判斷句 2 和句 1 之間的句子順序來學習文本一致性。
ALBERT 的設計相當成功,成為了壓縮 BERT 的經典范例,而為了做到最極限壓縮 BERT,以 ALBERT 作為起點確實是個不錯的主意。ALBERT 已經如此強大,EdgeBERT 又能壓到什么程度?作者一上來就用一張在 QQP 上的內存占用 / 運算量 / 性能比較圖吊足了讀者的胃口。(注意:縱坐標的內存占用量是對數刻度!)
本文對于 ALBERT 的利用除了作為初始化參數外,還在對下游任務做 fine-tune 時利用已經 fine-tune 好的 ALBERT 作為 teacher 進行知識蒸餾,來進一步提升模型性能。
初級配方:算法優化
- 基于熵的提前退出機制
出處: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 向 n 層的 BERT 模型添加了 (n-1) 個 “出口層” 分類器(Early Exit Off-Ramps)。出口層分類器 f_i 被放置在第 i 和 i+1 層 Transformer 之間,作為判斷第 i 層 Transformer 的信息是否足以進行推斷的標志。進行推斷時,從最底層的分類器開始逐層計算出口層分類器的熵,當某一層的出口層分類器熵小于某個閾值時,則將該出口層分類器結果作為模型結果,省去后續層的計算。
上圖表現的是不同的熵閾值在 MNLI、QQP、SST-2、QNLI 四個數據集上的平均退出層數、理論運行時間節省和相應的準確度。添加提前退出機制后,在 Acc 損失 1 個百分點時,能夠在這四個數據集上分別降低 30%、45%、54% 和 36% 的理論運行時間;在 Acc 損失 5 個百分點時,能將在這四個數據集上的理論運行時間的降低進一步降低至 44%、62%、78% 和 53%。
- 動態注意力范圍
出處: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 計算注意力,從而降低矩陣運算的開銷。
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 的參數。
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% 的準確度。
- 動態浮點數量化
出處: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 數據類型需要涉及到一些與機器學習無關的知識。
上圖的結果中,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 究竟能帶來多大的影響?上圖結果表明,僅僅是改變了嵌入層的硬件介質就能帶來約的推理延遲降低,以及低至約的能耗降低!這對于邊緣計算場景下簡直是質變了!(為何 ReRAM 只有讀,但 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 子結構依然是一個不錯的課題,至少它能讓更多人、更多時候、更多場景能夠用上性能強大的預訓練模型。本文中提到的這些優化方法是不是也給愛思考的你帶來了什么啟發呢?
https://mp.weixin.qq.com/s?__biz=MzA3MzI4MjgzMw==&mid=2650806805&idx=4&sn=4a356904e931334d3eaa7cbb51f3b69a&chksm=84e5de6bb392577d5a73c726590caabf034900e90e709c80eccc6e53e128a332e1808df58226&mpshare=1&scene=1&srcid=0114fb2F96KwO56sJC4IJOBt&sharer_sharetime=1610673987490&sharer_shareid=ad938198e7c1453de3019f515c855c6a&key=ff69355afd56518a3247a1644149d227b8bce5c9df9dbe51f3b17e614c710fac03f83e9d4ccda0cc910b7cb23d446c7d03e6b6b03a95bd2c181df8b481815762dd321fa5a0e3caace32de430594cbaaaa97128abaeba4ca9ab9e72325f67f23d5c8696f55643f21399d8553f733e1e00b676cd6d511f746745d74d667d7dfdb9&ascene=1&uin=MTk0ODYwODI2Mg%3D%3D&devicetype=Windows+10+x64&version=6300002f&lang=zh_CN&exportkey=Ae37qHjYq%2BzH3jTZVIfrUkw%3D&pass_ticket=xU9nO5ZDg9EM3GAMerffD%2BpFqLzXLQ3vDTL5jLQqhn86Pfu7VcD%2B0IvKP8QOaPNM&wx_header=0
總結
以上是生活随笔為你收集整理的EdgeBERT:极限压缩bert的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 样本不平衡不均衡数据处理
- 下一篇: java转python