文 | Cheney
編 | 智商掉了一地

既讓模型跑得快又好，又要空間占用少，真的有這樣效果爆表的操作存在嗎？

在當(dāng)前疫情反復(fù)的情況下，大家平時出門用健康寶刷臉的頻率變得越來越高。如果每次人臉識別都需要等很久的話，那也太讓人焦慮啦，甚至?xí)性沂謾C的沖動。而且我們還想讓模型的識別效果好，不能經(jīng)常出錯，就得用一些跑起來很慢的大模型來提升效果。那么問題來了，咋在這兩者之間做個平衡呢，既讓模型跑得又快又好，又要占用空間夠少，還能效果爆表？那就不得不提模型壓縮和最近興起的自動機器學(xué)習(xí)技術(shù)了，在壓縮模型的同時還能自動搜索出效果更優(yōu)的結(jié)構(gòu)，豈不美哉。

在NLP領(lǐng)域，預(yù)訓(xùn)練語言模型（如BERT等）在大量任務(wù)上都取得了不錯的效果，但是它們有著過大的體積以及推理時間，這嚴重限制了實際的部署。為了解決這個問題，很多模型壓縮方法被提出，例如知識蒸餾（Knowledge Distillation，KD）技術(shù)。雖然用于預(yù)訓(xùn)練模型壓縮的知識蒸餾方向取得了很大的進展，但針對Transformer結(jié)構(gòu)中的前饋神經(jīng)網(wǎng)絡(luò)（Feed-forward Network, FFN）壓縮和改進方面的研究還很少，其計算代價是多頭注意力（Multi-head Attention，MHA）模塊的2-3倍；同時[1]也指出，前饋網(wǎng)絡(luò)里的多層感知器（Multilayer Perceptron，MLP）結(jié)構(gòu)可以防止自注意機制帶來的秩崩塌（rank collapse）的問題，可以幫助模型的優(yōu)化。因此，如何在最大限度壓縮前饋網(wǎng)絡(luò)的同時，尋找到具有更優(yōu)秀非線性的前饋網(wǎng)絡(luò)結(jié)構(gòu)，進而提升模型效果，是個很值得關(guān)注的問題。

此外，自動機器學(xué)習(xí)技術(shù)，特別是模型架構(gòu)自動化搜索（Neural Architecture Search，NAS），在計算機視覺和自然語言處理領(lǐng)域已經(jīng)嶄露頭角，可以自動地從龐大的搜索空間中搜索出比人工設(shè)計的結(jié)構(gòu)效果更優(yōu)的模型。因此可以很自然地想到利用這類技術(shù)搜索出符合需求的前饋網(wǎng)絡(luò)結(jié)構(gòu)。但由于搜索空間往往極其龐大和復(fù)雜，如何高效地進行模型架構(gòu)搜索也是亟待解決的問題。

基于網(wǎng)絡(luò)架構(gòu)搜索的模型壓縮方法之前也有很多人進行研究，比如AdaBERT[2]基于卷積神經(jīng)網(wǎng)絡(luò)搜索任務(wù)自適應(yīng)的小模型。然而，這些方法沒有針對Transformer中前饋網(wǎng)絡(luò)的非線性進行探究；搜索空間受限于人工設(shè)計，不夠靈活；泛化性也值得考究。

為了解決上述問題，今天我們介紹一篇收錄于Findings of EMNLP 2021的工作EfficientBERT。在搜索空間方面，該方法主要針對Transformer中前饋網(wǎng)絡(luò)的數(shù)學(xué)表達式進行搜索，以發(fā)現(xiàn)具有更優(yōu)非線性的模型結(jié)構(gòu)；同時也將網(wǎng)絡(luò)的堆疊層數(shù)和中間通道數(shù)加入進來，以綜合平衡模型的參數(shù)量和推理時間。在搜索方法方面，該文提出了一種由粗到細的搜索策略，漸進地對搜索空間中的每個部分進行搜索，以提升搜索效率；同時提出了一種熱啟動知識蒸餾（Warm-up KD）策略，使每個搜索子結(jié)構(gòu)在訓(xùn)練時可以更快地收斂，并提升模型的泛化性能。

論文標(biāo)題
EfficientBERT: Progressively Searching Multilayer Perceptron via Warm-up Knowledge Distillation

論文鏈接
https://arxiv.org/abs/2109.07222

開源代碼
https://github.com/cheneydon/efficient-bert

1 方法

本文提出的方法主要分成兩個部分：搜索空間的設(shè)計，以及基于熱啟動知識蒸餾的漸進搜索方法。

▲EfficientBERT搜索整體框架

1.1 搜索空間

搜索空間設(shè)計方面，首先對前饋網(wǎng)絡(luò)的數(shù)學(xué)表達式進行搜索，以此反映網(wǎng)絡(luò)的非線性能力，并為其定義了如下幾種基本的數(shù)學(xué)操作符號：

▲候選數(shù)學(xué)操作符號

同時該文也對前饋網(wǎng)絡(luò)的堆疊層數(shù)和中間通道擴增比例進行搜索，以平衡模型的參數(shù)量和推理時間，分別從{1, 2, 3, 4}和{1, 1/2, 1/3, 1/4}中選取。在搜索時，每個前饋網(wǎng)絡(luò)結(jié)構(gòu)通過一個有向無環(huán)圖構(gòu)建，且上述操作被放入中間節(jié)點中來處理輸入特征。

值得一提的是，正如上述框架圖所示，該文對每層前饋網(wǎng)絡(luò)整體的數(shù)學(xué)表達式進行搜索，包括激活函數(shù)以及一些復(fù)雜的嵌套或組合表達，且各層網(wǎng)絡(luò)的表達式也各不相同。比如該文搜到的EfficientBERT模型第4, 5層的表達式為linear2_2(gelu(linear2_2(max(x,tanh(x)),wb1)),wb2)和linear1_2(relu(mul(linear1_2(x,wb1),x)),wb2)，其中l(wèi)inear x_y表示其所在的前饋網(wǎng)絡(luò)的堆疊層數(shù)為x、中間通道擴增比例為1/y，wb i表示第i個線性層的權(quán)重和偏置。

1.2 搜索方法

首先是基礎(chǔ)搜索模型（Base Model）結(jié)構(gòu)的設(shè)計。該文綜合采用了之前工作的一些壓縮方法，比如嵌入層因子分解（embedding factorization），減少模型寬度、深度、中間通道擴增比例等。

接下來是搜索過程，整個搜索過程被分為三個階段：

• 在第一階段，對整個搜索空間進行聯(lián)合搜索，其基于上述基礎(chǔ)搜索模型結(jié)構(gòu)。每個搜索子模型通過LaNAS[3]中提出的一種可學(xué)習(xí)的采樣決策樹進行采樣，使得更具有潛力的子模型可以以更大的概率被采樣到。為了避免不同數(shù)學(xué)表達式之間參數(shù)干擾的問題，在這個階段不同的子模型被單獨地進行訓(xùn)練。同時，為了使每個子模型的訓(xùn)練更快地收斂，該文提出一種熱啟動知識蒸餾方法。首先建立一個熱啟動超級網(wǎng)絡(luò)，該網(wǎng)絡(luò)有著最大的堆疊層數(shù)和中間通道擴增比例。將該超級網(wǎng)絡(luò)通過知識蒸餾預(yù)訓(xùn)練好之后固定其權(quán)重，在搜索時每個子模型從對應(yīng)位置繼承其權(quán)重，之后利用知識蒸餾只需要少量的預(yù)訓(xùn)練和微調(diào)步數(shù)即可達到收斂，節(jié)省了大量搜索時間。

• 在第二階段，對數(shù)學(xué)表達式進行單獨搜索，以發(fā)掘其更多樣化的表達并評估其對模型最終效果的影響，并將其它兩個部分的結(jié)構(gòu)固定住（即堆疊層數(shù)和中間通道擴增比例）。該階段基于第一階段搜索到的結(jié)構(gòu)進行搜索，并采用和第一階段相同的知識蒸餾方式。

• 在第三階段，對堆疊層數(shù)和中間通道擴增比例進行聯(lián)合搜索。該階段基于第二階段搜索到的結(jié)構(gòu)進行搜索，其數(shù)學(xué)表達式被固定，但使用了不同的堆疊層數(shù)和中間通道擴增比例組合，以增加搜索多樣性。由于無需對數(shù)學(xué)表達式進行搜索，該階段采用權(quán)重共享方法加速搜索。在該階段的熱啟動知識蒸餾過程中，將熱啟動超級網(wǎng)絡(luò)預(yù)訓(xùn)練好之后，其權(quán)重不再被固定；在搜索時，每個搜索子模型被均勻采樣，繼承超級網(wǎng)絡(luò)參數(shù)后進行訓(xùn)練，且不同子模型的參數(shù)可以共享。為了提高模型的泛化性，在訓(xùn)練時采用了多任務(wù)訓(xùn)練的方法，讓每個子模型在不同種類的下游任務(wù)上進行微調(diào)。其中嵌入層和Transformer層的參數(shù)在所有任務(wù)上共享，但不對預(yù)測層參數(shù)進行共享。

2 實驗

首先是在GLUE test和SQuAD dev數(shù)據(jù)集上與之前的模型壓縮方法進行比較：

▲GLUE test數(shù)據(jù)集比較結(jié)果▲SQUAD dev數(shù)據(jù)集比較結(jié)果

并在GLUE dev數(shù)據(jù)集上與之前相關(guān)的利用架構(gòu)搜索進行模型壓縮的方法進行比較：

▲GLUE dev數(shù)據(jù)集比較結(jié)果

可以看到，在相似的參數(shù)量下，本文搜索到的模型EfficientBERT效果遠遠超過了之前效果最好的模型壓縮方法（TinyBERT[4], MobileBERT[5], MiniLM[6]等），也超過了之前利用架構(gòu)搜索的模型壓縮方法（AdaBERT[2]，NAS-BERT[7]等）。

接下來，本文也給出各階段搜索到模型的最終性能，驗證提出的漸進架構(gòu)搜索方法、以及搜索空間中各個部分的有效性：

▲不同搜索階段模型最終性能

本文也給出各個階段的搜索時間，驗證提出的漸進搜索方法的加速能力：

▲由粗到細搜索方法各階段搜索時間

可以看到，本文提出的漸進搜索方法不僅可以搜索出更具有潛力的模型結(jié)構(gòu)，而且可以加速搜索進程。

為了測試搜到的模型結(jié)構(gòu)的泛化能力，EfficientBERT結(jié)構(gòu)隨后被遷移到了TinyBERT上：

▲模型結(jié)構(gòu)遷移效果

此外，本文也給出各個子模型在搜索階段和最終測試階段預(yù)測結(jié)果的相關(guān)性結(jié)果，驗證所提出的熱啟動知識蒸餾方法對不同子模型效果的預(yù)測排序能力：

▲熱啟動知識蒸餾排序能力

可以看到，使用熱啟動知識蒸餾方法，子模型在各個下游任務(wù)上都可以在搜索和最終測試階段的預(yù)測結(jié)果之間保持很高的正相關(guān)性。

為了直觀地評估模型的非線性能力，本文對搜索到的前饋網(wǎng)絡(luò)結(jié)構(gòu)進行了可視化：

▲前饋網(wǎng)絡(luò)非線性能力可視化

其中(d)-(f)是隨機采樣出的效果更差的幾個搜索子模型。從上圖可以看出，(a)-(c)相比于(d)-(f)曲線更加流暢且異常突增區(qū)域更少；并且從(a)到(c)，曲線的復(fù)雜程度越來越低，顯示了搜索到的模型EfficientBERT優(yōu)秀的非線性能力。

最后，本文將各個階段搜索到的模型整體結(jié)構(gòu)進行了可視化：

▲模型整體結(jié)構(gòu)圖

從最終階段搜到的EfficientBERT結(jié)構(gòu)中可以發(fā)現(xiàn)一些有趣的規(guī)律：

• 由于多數(shù)前饋網(wǎng)絡(luò)中間通道擴增比例為1/2，且多數(shù)前饋網(wǎng)絡(luò)堆疊層數(shù)少于2，因此搜索得到的EfficientBERT非常輕量；

• 更低層有著更大的堆疊層數(shù)或中間通道擴增比例（如第1, 2層），以最大程度地豐富模型的語義表達能力；

• 更高層有著更復(fù)雜的數(shù)學(xué)表達式（如第4, 5層），以提高模型的非線性表達能力。

3 總結(jié)

本文主要針對前饋網(wǎng)絡(luò)的尺寸和非線性能力進行壓縮和提升，設(shè)計了一種非常靈活龐大的搜索空間進行自動化架構(gòu)搜索。為了提高搜索效率，本文提出了一種由粗到細的架構(gòu)搜索方法，并在每個搜索階段采用了一種熱啟動知識蒸餾方法來加速搜索并提高模型的泛化能力。大量實驗證明了搜索到的模型EfficientBERT的有效性。

后臺回復(fù)關(guān)鍵詞【入群】

加入賣萌屋NLP/IR/Rec與求職討論群

后臺回復(fù)關(guān)鍵詞【頂會】

獲取ACL、CIKM等各大頂會論文集！

?

[1] Yihe Dong, et al. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. ICML 2021.

[2] Daoyuan Chen, et al. AdaBERT: Task-adaptive bert compression with differentiable neural architecture search. IJCAI 2020.

[3] Linnan Wang, et al. Sample-efficient neural architecture search by learning action space. TPAMI 2021.

[4] Xiaoqi Jiao, et al. TinyBERT: Distilling BERT for natural language understanding. Findings of EMNLP 2020.

[5] Zhiqing Sun, et al. MobileBERT: a compact task-agnostic BERT for resource-limited devices. ACL 2020.

[6] Wenhui Wang, et al. MiniLM: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. NeurIPS 2020.

[7] Jin Xu, et al. NAS-BERT: Task-agnostic and adaptive-size bert compression with neural architecture search. KDD 2021.

總結(jié)

以上是生活随笔為你收集整理的EMNLP'21 | 让压缩语言模型自动搜索最优结构！的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。