EMNLP 2020 | 忒修斯之船啟發下的知識蒸餾新思路

微軟研究院AI頭條發表於2020-11-11
隨著深度學習的流行,很多大型神經網路模型誕生,並在多個領域中取得當前最優的效能。尤其是在自然語言處理(NLP)領域中,預訓練和調參已經成為其中大多數任務的新正規化。基於 Transformer 的預訓練模型在自然語言理解(NLU)和自然語言生成(NLG)領域中成為主流。儘管這些模型從“過引數化”的特性中獲益,但它們往往包含數百萬甚至數十億個引數,這就使得此類模型的計算成本高昂,且從記憶體消耗和高延遲的角度來看計算低效。這一缺陷極大地阻礙了此類模型在生產環境中的應用。

為了解決該問題,研究人員提出了很多神經網路壓縮技術。一般而言,這些技術可以分為三類:量化權重剪枝知識蒸餾(Knowledge Distillation)。其中,由於知識蒸餾能夠壓縮預訓練語言模型,所以得到了極大關注。知識蒸餾利用大型教師模型“教”緊湊的學生模型模仿教師的行為,從而將教師模型中嵌入的知識遷移到較小的模型中。但是,學生模型的效能狀況取決於設計良好的蒸餾損失函式,正是這個函式使得學生模型可以模仿教師的行為。近期關於知識蒸餾的研究甚至利用更復雜的模型特定蒸餾損失函式,以實現更好的效能。

近日,來自微軟亞洲研究院自然語言計算組的研究員們提出了一種與顯式地利用蒸餾損失函式來最小化教師模型與學生模型距離的知識蒸餾不同的模型壓縮新方法。受到著名哲學思想實驗“忒修斯之船”的啟發(即如果船上的木頭逐漸被替換,直到所有的木頭都不是原來的木頭,那這艘船還是原來的那艘船嗎?),研究員們在 EMNLP 2020 上發表了 Theseus Compression for BERT (BERT-of-Theseus),該方法逐步將 BERT 的原始模組替換成引數更少的替代模組(點選文末閱讀原文,瞭解論文詳情)。研究員們將原始模型叫做“前輩”(predecessor),將壓縮後的模型叫做“接替者”(successor),分別對應知識蒸餾中的教師和學生。

該方法的工作流程如下圖所示。首先為每個前輩模組指定一個接替者模組,然後在訓練階段中以一定的機率(如拋硬幣)決定是否用替代模組隨機替換對應的前輩模組,並按照新舊模組組合的方式繼續訓練。在模型收斂後,將所有接替者模組組合成接替者模型,進而執行推斷。這樣就可以將大型前輩模型壓縮成緊湊的接替者模型了。

EMNLP 2020 | 忒修斯之船啟發下的知識蒸餾新思路

舉例來說,假設現在有兩支籃球隊每支各五人,一支是經驗老道的全明星球隊,另一支則是年輕球員組成的青訓隊。為了提高青訓隊的水平,所以隨機選派青訓隊員去替換掉全明星隊中的球員,然後讓這個混合的球隊不斷地練習、比賽。透過向前輩學習經驗,新加入成員的實力會有所提升,也能學會和其他隊員的配合,逐漸的這個混合球隊就擁有了接近全明星球隊的實力。之後重複這個過程,直到青訓隊員都被充分訓練,最終青訓隊員也能自己組成一支實力突出的球隊。相比之下,如果沒有“老司機”來帶一帶,青訓隊無論如何訓練,水平也不會達到全明星隊的實力。

事實上,Theseus 壓縮與知識蒸餾的思路有些類似,都是鼓勵壓縮模型模仿原始模型的行為,但 Theseus 壓縮有很多獨特的優勢。

首先,Theseus 壓縮在壓縮過程中僅使用任務特定的損失函式。而基於知識蒸餾的方法除了使用任務特定的損失函式外,還需加入繁瑣的蒸餾損失函式作為最佳化目標。

其次,與近期研究 TinyBERT 等不同,Theseus 壓縮不使用 Transformer 特定特徵進行壓縮,這就為壓縮廣泛模型提供了可能性。與知識蒸餾僅使用原始模型執行推斷不同,該方法允許前輩模型與壓縮後的接替者模型共同訓練,從而實現更深層次的梯度級互動,並簡化訓練過程。

此外,混合了前輩模組和接替者模組的不同模組組合還新增了額外的正則化項(類似於 Dropout)。該方法基於課程學習(Curriculum Learning)方法來驅動模組替換,將模組替換機率從低到高逐漸增加,從而實現優異的 BERT 壓縮效能。利用Theseus 壓縮方法壓縮得到的 BERT 模型運算速度是之前的1.94 倍,並且保留了原始模型超過 98% 的效能,優於其它基於知識蒸餾的壓縮的基線方法。

EMNLP 2020 | 忒修斯之船啟發下的知識蒸餾新思路

透過在預訓練語言模型 BERT 上的成功實驗,微軟亞洲研究院的研究員們希望可以為模型壓縮開啟一種全新的思路,並希望看到這一方法在計算機視覺等領域的更廣泛應用。

相關文章