文 | Cheney
編 | 智商掉了一地

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

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

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

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

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

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

論文標題
EfficientBERT: Progressively Searching Multilayer Perceptron via Warm-up Knowledge Distillation

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

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

1 方法

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

▲EfficientBERT搜索整體框架

1.1 搜索空間

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

▲候選數學操作符號

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

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

1.2 搜索方法

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

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

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

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

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

2 實驗

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

▲GLUE test數據集比較結果▲SQUAD dev數據集比較結果

并在GLUE dev數據集上與之前相關的利用架構搜索進行模型壓縮的方法進行比較：

▲GLUE dev數據集比較結果

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

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

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

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

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

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

為了測試搜到的模型結構的泛化能力，EfficientBERT結構隨后被遷移到了TinyBERT上：

▲模型結構遷移效果

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

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

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

為了直觀地評估模型的非線性能力，本文對搜索到的前饋網絡結構進行了可視化：

▲前饋網絡非線性能力可視化

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

最后，本文將各個階段搜索到的模型整體結構進行了可視化：

▲模型整體結構圖

從最終階段搜到的EfficientBERT結構中可以發現一些有趣的規律：

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

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

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

3 總結

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

[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.

往期精彩回顧適合初學者入門人工智能的路線及資料下載機器學習及深度學習筆記等資料打印機器學習在線手冊深度學習筆記專輯《統計學習方法》的代碼復現專輯 AI基礎下載黃海廣老師《機器學習課程》視頻課黃海廣老師《機器學習課程》711頁完整版課件

本站qq群851320808，加入微信群請掃碼：

與50位技術專家面對面20年技術見證，附贈技術全景圖

總結

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

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