?PaperWeekly 原創 ·?作者｜蘇劍林

學校｜追一科技

研究方向｜NLP、神經網絡

最近了解到一種稱為“BERT-of-Theseus”的 BERT 模型壓縮方法，來自論文?BERT-of-Theseus: Compressing BERT by Progressive Module Replacing。這是一種以“可替換性”為出發點所構建的模型壓縮方案，相比常規的剪枝、蒸餾等手段，它整個流程顯得更為優雅、簡潔。

論文標題：BERT-of-Theseus: Compressing BERT by Progressive Module Replacing

論文鏈接：https://arxiv.org/abs/2002.02925

本文將對該方法做一個簡要的介紹，給出一個基于 bert4keras [1] 的實現，并驗證它的有效性。

▲ BERT-of-Theseus，原作配圖

模型壓縮

首先，我們簡要介紹一下模型壓縮。不過由于筆者并非專門做模型壓縮的，也沒有經過特別系統的調研，所以該介紹可能顯得不專業，請讀者理解。

1.1 基本概念

簡單來說，模型壓縮就是“簡化大模型，得到推理速度更快的小模型”。當然，一般來說模型壓縮是有一定犧牲的，比如最明顯的是最后的評測指標會有一定的下降，畢竟“更好又更快”的免費午餐是很少的，所以選擇模型壓縮的前提是能允許一定的精度損失。

其次，模型壓縮的提速通常只體現在預測階段，換句話說，它通常需要花費更長的訓練時間，所以如果你的瓶頸是訓練時間，那么模型壓縮也不適合你。

模型壓縮要花費更長時間的原因是它需要“先訓練大模型，再壓縮為小模型”。讀者可能會疑惑：為什么不直接訓練一個小模型？答案是目前很多實驗已經表明，先訓練大模型再壓縮，相比直接訓練一個小模型，最后的精度通常會更高一些。

也就是說，在推理速度一樣的情況，壓縮得到的模型更優一些，相關探討可以參考論文 Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers? [2]?，另外知乎上也有討論《為什么要壓縮模型，而不是直接訓練一個小的 CNN？》[3]。

1.2 常見手段

常見的模型壓縮技術可以分為兩大類：1）直接簡化大模型得到小模型；2）借助大模型重新訓練小模型。這兩種手段的共同點是都先要訓練出一個效果比較好的大模型，然后再做后續操作。

第一類的代表方法是剪枝（Pruning）和量化（Quantization）。

剪枝，顧名思義，就是試圖刪減掉原來大模型的一些組件，使其變為一個小模型，同時使得模型效果在可接受的范圍內；

至于量化，指的是不改變原模型結構，但將模型換一種數值格式，同時也不嚴重降低效果，通常我們建立和訓練模型用的是 float32 類型，而換成 float16 就能提速且省顯存，如果能進一步轉換成 8 位整數甚至 2 位整數（二值化），那么提速省顯存的效果將會更加明顯。

第二類的代表方法是蒸餾（Distillation）。蒸餾的基本想法是將大模型的輸出當作小模型訓練時的標簽來用，以分類問題為例，實際的標簽是 one hot 形式的，大模型的輸出（比如 logits）則包含更豐富的信號，所以小模型能從中學習到更好的特征。

除了學習大模型的輸出之外，很多時候為了更進一步提升效果，還需要小模型學習大模型的中間層結果、Attention 矩陣、相關矩陣等，所以一個好的蒸餾過程通常涉及到多項 loss，如何合理地設計這些 loss 以及調整這些 loss 的權重，是蒸餾領域的研究主題之一。

Theseus

本文將要介紹的壓縮方法稱為“BERT-of-Theseus”，屬于上面說的兩大類壓縮方法的第二類，也就是說它也是借助大模型來訓練小模型，只不過它是基于模塊的可替換性來設計的。

BERT-of-Theseus 的命名源于思想實驗“忒修斯之船”：如果忒修斯的船上的木頭被逐漸替換，直到所有的木頭都不是原來的木頭，那這艘船還是原來的那艘船嗎？

2.1 核心思想

前面說到，用蒸餾做模型壓縮時，往往不僅希望小模型的輸出跟大模型的輸出對齊，還希望中間層結果也對齊。“對齊”意味著什么呢？意味著可替換！

所以 BERT-of-Theseus 的思想就是：干嘛要煞費苦心地通過添加各種 loss 去實現可替換性呢？直接用小模型的模塊去替換掉大模型的模塊然后去訓練不就好了嗎？

舉個實際的類比：

假設現在有 A、B 兩支球隊，每支各五人。A 球隊屬于明星球隊，實力超群；B 球隊則是新手球隊，待訓練。為了訓練 B 球隊，我們從 B 球隊中選 1 人，替換掉 A 球隊中的 1 人，然后讓這個“4+1”的A球隊不斷的練習、比賽。經過一段時間，新加入的成員實體會提升，這個“4+1”的球隊就擁有接近原始 A 球隊的實力。

重復這個過程，直到 B 球隊的人都被充分訓練，那么最終 B 球隊的人也能自己組成一支實力突出的球隊。相比之下， 如果一開始就只有 B 球隊，只是 B 球隊的人自己訓練、比賽，那么就算他們的實力逐漸提升，但由于沒有實力超群的 A 球隊幫助，其最終實力也不一定能突出。

2.2 流程細節

回到BERT的壓縮，現在假設我們有一個 6 層的 BERT，我們直接用它在下游任務上微調，得到一個效果還不錯的模型，我們稱之為?Predecessor（前輩）。

我們的目的是得到一個 3 層的 BERT，它在下游任務重效果接近 Predecessor，至少比直接拿 BERT 的前 3 層去微調要好（否則就白費力氣了），這個小模型我們稱為?Successor（后輩）。那么 BERT-of-Theseus 是怎么實現這一點的呢？如下圖。

▲ BERT-of-Theseus訓練過程示意圖

▲?Predecessor和Successor模型示意圖

在 BERT-of-Theseus 的整個流程中，Predecessor 的權重都被固定住。6 層的 Predecessor 被分為 3 個模塊，跟 Successor 的 3 層模型一一對應，訓練的時候，隨機用 Successor 層替換掉 Predecessor 的對應模塊，然后直接用下游任務的優化目標進行微調（只訓練 Successor 的層）。

訓練充分后，再把整個 Successor 單獨分離出來，繼續在下游任務中微調一會，直到驗證集指標不再上升。

▲ 上述模型的等效模型

在實現的時候，事實上是類似 Dropout 的過程，同時執行 Predecessor 和 Successor 模型，并將兩者對應模塊的輸出之一置零，然后求和、送入下一層中，即：

由于 非 0 即 1（不作調整，各自 0.5 概率隨機選效果就挺好了），所以每個分支其實就相當于只有一個模塊被選擇到，因此上面右圖就相當于下述模型結構。由此每次的置零都是隨機的，因此訓練足夠多的步數后，Successor 的每個層都能被訓練好。

2.3 方法分析

跟蒸餾相比，BERT-of-Theseus 有什么優勢呢？首先，這既然能被發表出來，所以至少效果應該是不相上下的，所以我們就不去比較效果了，而是比較方法本身。很明顯，BERT-of-Theseus 的主要特點是：簡潔。

前面說到，蒸餾多數時候也需要匹配中間層輸出，這時候要涉及到的訓練目標就有很多了：下游任務 loss、中間層輸出 loss、相關矩陣 loss、Attention 矩陣 loss、等等，想想要平衡這些 loss 就是一件頭疼的事情。

相比之下，BERT-of-Theseus 直接通過替換這個操作，逼著 Successor 能有跟 Predecessor 類似的輸出，而最終的訓練目標就只有下游任務 loss，不可謂不簡潔。

此外，BERT-of-Theseus 還有一個特別的優勢：很多的蒸餾方法都得同時作用于預訓練和微調階段，效果才比較突出，而 BERT-of-Theseus 直接作用于下游任務的微調，就可以得到相媲美的效果。這個優勢在算法上體現不出來，屬于實驗結論。

從形式上來看，BERT-of-Theseus 的隨機替換思路有點像圖像中的數據擴增方案 SamplePairing 和 mixup（參考《從 SamplePairing到mixup：神奇的正則項》[4] ），都是隨機采樣兩個對象加權求和來增強原模型；也有點像 PGGAN [5] 的漸進式訓練方案，都是通過對兩個模型進行某種程度的混合，實現兩個模型的過渡。

如果了解它們的讀者，繼而就能夠對 BERT-of-Theseus 提出一些拓展或者說疑問： 一定要非 0 即 1 嗎，任意 的隨機數行不？或者說不隨機，直接讓 慢慢地從 1 變到 0 行不？這些想法都還沒有經過充分實驗，有興趣的讀者可以修改下述代碼自行實驗。

實驗效果

原作者們開源了自己的 PyTroch 實現：

https://github.com/JetRunner/BERT-of-Theseus

邱震宇老兄也分享了自己的講解 [6] 以及基于原版 BERT 的 Tensorflow 實現 qiufengyuyi/bert-of-theseus-tf [7]。當然，既然筆者決定寫這篇介紹，那就肯定少不了基于 bert4keras 的 Keras 實現了：

https://github.com/bojone/bert-of-theseus

這大概是目前最簡潔、最具可讀性的 BERT-of-Theseus 實現了，沒有之一。

原論文的效果大家就自己去看原論文了。筆者在幾個文本分類任務上實驗了一下，結果大同小異，跟邱兄的實驗結論也比較一致。其中在 CLUE 的 iflytek 數據集中實驗結果如下：

可以看到，相比直接拿前幾層微調，BERT-of-Theseus 確實能帶來一定的性能提升。對于隨機置零方案，除了均等概率選擇 0/1 外，原論文還嘗試了其他策略，有輕微提升，但會引入額外超參，所以筆者就沒有實驗了，有興趣的讀者可以自己修改嘗試。

另外，對于蒸餾來說，如果 Successor 跟 Predecessor 有同樣的結構（同模型蒸餾），那么通常來說 Successor 的最終性能比 Predecessor 還要好些，BERT-of-Theseus 有沒有這一特點呢？

筆者也實驗了一下該想法，發現結論是否定的，也就是同模型情況下 BERT-of-Theseus 訓練出來的 Successor 并沒有比 Predecessor 好，所以看來 BERT-of-Theseus 雖好，但也不能完全取代蒸餾。

文末小結

本文介紹并實驗了一種稱為“BERT-of-Theseus”的 BERT 模型壓縮方法，該方法的特點是簡潔明了，純粹通過替換操作來讓小模型去學習大模型的行為，使得能在只有一個 loss 的情況下就能達到當前最優的模型壓縮效果。

參考文獻

[1] https://github.com/bojone/bert4keras

[2] https://arxiv.org/abs/2002.11794

[3] https://www.zhihu.com/question/303922732

[4] https://kexue.fm/archives/5693

[5] https://arxiv.org/abs/1710.10196

[6] https://zhuanlan.zhihu.com/p/112787764

[7] https://github.com/qiufengyuyi/bert-of-theseus-tf

更多閱讀

#投 稿?通 道#

?讓你的論文被更多人看到?

如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學習心得或技術干貨。我們的目的只有一個，讓知識真正流動起來。

?????來稿標準：

? 稿件確系個人原創作品，來稿需注明作者個人信息（姓名+學校/工作單位+學歷/職位+研究方向）?

? 如果文章并非首發，請在投稿時提醒并附上所有已發布鏈接?

? PaperWeekly 默認每篇文章都是首發，均會添加“原創”標志

?????投稿郵箱：

? 投稿郵箱：hr@paperweekly.site?

? 所有文章配圖，請單獨在附件中發送?

? 請留下即時聯系方式（微信或手機），以便我們在編輯發布時和作者溝通

????

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

關于PaperWeekly

PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

總結

以上是生活随笔為你收集整理的BERT-of-Theseus：基于模块替换的模型压缩方法的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。