編者按：如今，基于 Transformer 的大規(guī)模預訓練語言模型，如 BERT、XLNE、RoBERTa 和 GPT-3 等，已經(jīng)在很多自然語言處理任務中都取得了十分驚人的效果。但是巨大的模型尺寸，使其在眾多不同的下游任務中進行部署時非常困難。而且由于存在大量復雜的場景以及不同的下游任務，單獨為不同場景設計一種壓縮過的 BERT 模型既耗時又耗力。

因此，微軟亞洲研究院的研究員們針對這些問題展開了研究，并提出了 NAS-BERT 技術。相關研究論文“NAS-BERT：Task-Agnostic and Adaptive-Size BERT Compression with Neural Architechture Search”已被跨學科數(shù)據(jù)科學會議 KDD 2021 收錄。歡迎感興趣的讀者點擊閱讀原文，查看論文全文。

不同設備往往對訓練的模型有不同的要求，如模型占用內存的大小、推理時延等。為了滿足設備的要求，技術人員往往需要對大規(guī)模的預訓練模型進行壓縮處理。但是由于存在大量復雜的場景和不同的下游任務，專門為每一種場景設計壓縮過的BERT模型，需要投入大量的人力和計算資源。

對此，微軟亞洲研究院的研究員們提出了 NAS-BERT 技術（相關論文已發(fā)表在 KDD 2021 大會上），可以直接在上游預訓練任務中進行壓縮訓練，使壓縮模型不再依賴于下游的具體任務。并且 NAS-BERT 只需通過一次搜索就可以直接產生不同尺寸的模型，進而滿足不同設備的要求。所以即使是資源受限的設備，也可以動態(tài)選擇合適的架構。相較于傳統(tǒng)方法，NAS-BERT 極大地提高了訓練效率。

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

如表1所示，之前大部分的壓縮工作都是針對具體任務本身而設計方法實現(xiàn)的。比如，DisilBERT 等工作會直接在上游模型進行壓縮學習，對下游任務無關，但不能適應各種不同受限資源的場景。DynaBERT 則可以在下游任務中通過一次訓練產生多種不同大小的模型，以達到動態(tài)調整尺寸的目的。而 NAS-BERT 與其他幾種工作不同，它可以同時做到任務無關和動態(tài)調整尺寸，并且達到優(yōu)越的性能。

表1：之前的 BERT 壓縮工作

為了讓 NAS-BERT 實現(xiàn)上述目的，研究員們直接在上游預訓練階段對 NAS-BERT 進行了架構搜索（NAS）與訓練，并且一次性搜索出各種大小的模型，以方便各種場景的部署。通過嘗試架構層面的組合方式，NAS-BERT 不僅能夠探索模型本身的潛力，還可以搜索出更加優(yōu)越的架構。

具體而言，NAS-BERT 首先在上游預訓練任務上訓練一個超網(wǎng)絡。其中，超網(wǎng)絡中的一個有向無環(huán)圖就是一個架構，不同的架構在超網(wǎng)絡中共享參數(shù)來降低訓練成本。然后研究員們再通過自動架構搜索技術去訓練這個超網(wǎng)絡。由于超網(wǎng)絡中包括各種尺寸的模型，因此可以覆蓋到各種不同的模型大小。最后，在用超網(wǎng)絡去評價每個架構的性能并選擇出優(yōu)越的架構。

然而在 BERT 預訓練任務上做自動架構搜索是具有挑戰(zhàn)性的。其主要原因有：1. BERT 預訓練本身收斂慢并且需要巨大的計算資源；2. 巨大的搜索空間會進一步導致訓練超網(wǎng)絡變得困難。為了解決這兩個問題，NAS-BERT 主要采用了模塊化搜索和逐漸縮小搜索空間的技術。同時，為了在縮小搜索空間的基礎上能夠給出各種不同的架構，NAS-BERT 還使用了分桶搜索把搜索空間分成眾多不同的桶，不同桶中的架構有著不同的模型大小和延遲，桶中的架構相互競爭，從而選出最優(yōu)越的架構。并且不同桶的架構通過權重共享，以降低架構搜索的代價。通過使用這些技術，NAS-BERT 才得以快速地搜索出各種不同的優(yōu)越架構。

搜索空間的架構設計

NAS-BERT 的架構設計（如圖1所示），首先使用了一個預訓練的 BERT 模型作為教師模型，然后用教師模型指導超網(wǎng)絡的搜索。為了降低搜索空間，NAS-BERT 把學生模型（超網(wǎng)絡）分成了幾個不同的搜索模塊，然后每個搜索模塊會進行單獨的訓練。超網(wǎng)絡模塊的監(jiān)督信號由教師模型的相應模塊給出。換言之，就是給定一批數(shù)據(jù)，可以得到教師模型的每個模塊輸入輸出的隱藏表征，再用這些輸入輸出的隱藏表征去訓練搜索模塊。這樣，不同的搜索模塊可以分開搜索且并行訓練，指數(shù)級地降低了搜索空間。

對于每一個搜索模塊，網(wǎng)絡設計如圖1(a)所示，網(wǎng)絡的每一層包括了所有可能的搜索的候選操作（例如卷積、注意力網(wǎng)絡）。一個架構就是一條從底部到最上層的單向路徑，不同的架構（路徑）共享權重以降低訓練所需要的資源。

圖1：NAS-BERT 方法框架圖

為了使搜索的架構多樣化，研究員們分別采用了三種不同的操作方法：注意力網(wǎng)絡 MHA，前饋網(wǎng)絡 FFN 和可分離卷積 CONV。MHA 和 FFN 是 Transformer 模型中的操作方法，采用這兩種方法，可以使 Transformer 容納在搜索空間中。由于 CONV 已經(jīng)在很多自然語言任務中取得了不錯的效果，因此研究員們也將其加入了搜索空間，以探索卷積和其他操作組合的潛力（操作設計請見圖2）。同時，為了使搜索的架構具有各種不同的模型大小，研究員們讓每一種操作都有128/192/256/384/512這5種不同的隱藏層大小，如表2所示。與此同時，還加入了 Identity（無操作），使其能搜索出各種不同層的架構。

表2：操作集合的設計

圖2：卷積操作的設計

通過采用這種設計，可以使一個6層的搜索模塊中有超過20萬種可能的架構。如果采用4個搜索模塊，架構的組合方式將超過10^20種。為了訓練這個超網(wǎng)絡，每次得到一批數(shù)據(jù)，就要隨機采樣一個架構進行訓練，進而監(jiān)督和學習教師模塊的隱藏表征。由于搜索模塊采樣出來的架構輸入輸出的隱藏層大小可能和教師模塊不一致，所以可以直接引入一個可學習的線性層進行轉換，并和超網(wǎng)絡一起訓練。

漸進架構搜索，分桶設計與模型選擇

雖然分模塊搜索已經(jīng)大大減小了搜索空間，但是訓練超網(wǎng)絡依然需要很長的時間收斂。而且大量的模型在競爭有限的資源，將導致每個模型的評估結果都不準確，嚴重影響了架構搜索的排序性能。所以研究員們在訓練過程中的每一輪 epoch 結束時，會把搜索空間中一些沒有希望的架構裁剪掉，然后在裁剪后的搜索空間上繼續(xù)訓練。通過這種方式，可以將更多的計算資源給到更有希望的架構，進而得到更加準確的評估，這既減少了資源的浪費，又加快了搜索過程。

然而直接對整個搜索空間進行裁剪，會導致無法得到大量不同大小的架構。因為超網(wǎng)絡在搜索過程中可能偏向于大模型或者收斂快的模型，所以無約束的裁剪會使得最后得到的架構大小趨于相同?；诖?，研究員們提出了分桶裁剪的方法，并且采用具體的數(shù)字為例，闡述了在分桶裁剪過程中，搜索空間是如何變化的。

首先研究員們將搜索空間均勻地分成10個桶，每個桶的架構被約束在不同的參數(shù)量和延遲下。盡管不同的架構落在不同的桶里，但是在訓練的時候，它們的參數(shù)仍然是共享的。在每個迭代訓練結束的時候，研究員們會隨機從每個桶里采樣出2000個架構，同時使用超網(wǎng)絡評價它們在驗證集上的分數(shù)，并刪除一半的架構。通過重復這個過程，可以不斷地收縮搜索空間，直到每個桶里只剩下10個架構。這樣對于一個搜索模塊而言，最終就得到了100個架構。假設共有4個搜索模塊，那么就有100^4種組合方式，每一種組合方式都是一個完整的架構。

由于每個搜索模塊中的100個架構來自不同的桶，因此它們有著顯著不同的大小和延遲。所以通過組合不同的搜索模塊而得到的100^4種架構，也具有顯著不同的模型大小和延遲。

為了評估每個架構的性能，研究員們會首先評價每個搜索模塊中100個架構在驗證集上的損失。然后用不同模塊的損失直接加和作為組合架構的性能。通過這種方式，僅僅通過4*100次驗證集測試，就可以粗略得到所有架構的性能。該方法對于延遲的評估也是類似的，即先測量每個搜索模塊內架構的延遲，然后再用不同模塊的延遲加和作為整個架構的延遲。

所以當給定任何一種關于模型大小和延遲的約束時，可以通過查表的方式快速找到性能最高的架構。這種架構本身是在上游預訓練任務中挑選出來的，因此與下游任務無關。

實驗結果

為了評估搜索得到的架構性能，研究員們重新在上游任務上訓練了這個架構，然后在各種下游任務中測試了模型的性能。

在搜索階段使用的教師模型是 BERT 110M 的模型。教師模型和搜索出來的架構都是在 BookCorpus+English Wikipedia (16GB) 數(shù)據(jù)上進行了預訓練。為了和 BERT 的12層 Transformer 對齊，超網(wǎng)絡包含了24個子層（每一個 Transformer 層等于一層 MHA 加一層FFN）。研究員們在大量的下游任務上評估搜索出來的架構的性能，包括 GLUE 和 SQuAD 數(shù)據(jù)集，并選擇了5M，10M，30M，60M參數(shù)量的架構為例，來展示搜索出來的架構效果。

研究員們首先把 NAS-BERT 搜索出來的架構和手工設計的 BERT 架構進行了對比。從表3中可以看到，NAS-BERT 在各種模型大小及延遲條件下，都超過了手工設計的 BERT 模型。

表3：NAS-BERT 和 BERT 對比

進一步，研究員們將 NAS-BERT 和之前的模型壓縮工作進行對比。之前的模型壓縮工作主要采用了多種蒸餾技巧和訓練方式。而 NAS-BERT 只使用了簡單的兩階段蒸餾來突出其搜索出的架構優(yōu)勢，不使用復雜的蒸餾技術或訓練方法，如注意力蒸餾，逐層替換等。從表4可以看到，NAS-BERT 超過了之前的方法，證明了架構本身的優(yōu)越性。

表4：NAS-BERT 和之前的 BERT 壓縮工作對比

為了證明縮小搜索空間方法的有效性，研究員們又做了一組實驗——與不使用逐漸縮小搜索空間進行對比。從圖3的損失函數(shù)曲線來看，使用逐漸縮小搜索空間顯著加快了收斂速度。從最終搜索得到的架構來看，使用逐漸縮小搜索空間能夠幫助搜索到更好的架構。

圖3：逐漸縮小搜索空間的分離實驗

同時，研究員們還探究了縮小搜索空間的其他方法。與從架構層面縮小搜索空間不同，研究員們嘗試了從操作層面縮小搜索空間的方法（具體見論文）。從表5可以看到，從架構層面縮小搜索空間更為準確，因此能搜索到更好的架構。

表5：不同逐漸縮小搜索空間的分離試驗

最后，研究員們使用了各種不同的訓練方式來訓練 NAS-BERT，并將其與 BERT 模型對比，以證明搜索的架構的魯棒性。并且研究員們還嘗試了在訓練過程中，在上游預訓練階段或下游微調階段或兩者都有的情況下，觀察 NAS-BERT 搜索的架構的魯棒性。從表6中可以看到，NAS-BERT 在各種訓練配置上，都能顯著超過手工設計的 BERT 架構。

表6：不同訓練方式的分離實驗

表7展示了部分 NAS-BERT 搜索得到的架構?？梢园l(fā)現(xiàn)，搜索出來的架構都由不同的操作且復雜的方式組成，這證明了 NAS-BERT 能夠搜索出更加新穎的新架構。

表7：部分 NAS-BERT 搜索得到的架構

微軟亞洲研究院的研究員們在本篇論文中提出的 NAS-BERT，是一種用自動架構搜索方法實現(xiàn)任務無關且動態(tài)調整尺寸的 BERT 壓縮技術。其擁有新穎的搜索空間、卷積、注意力、前饋網(wǎng)絡以及不同的隱藏層大小。加上高效的搜索方法，NAS-BERT 可以探索出不同操作的復雜組合方式得到模型的潛力。研究員們通過大量的比較和分離實驗，也證明了 NAS-BERT 搜索得到架構的有效性。