KDD 2021 | 用NAS實(shí)現任務(wù)無(wú)關(guān)且可動(dòng)態(tài)調整尺寸的BERT壓縮（2）

漸進(jìn)架構搜索，分桶設計與模型選擇

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

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

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

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

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

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

實(shí)驗結果

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

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

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

表3：NAS-BERT 和 BERT 對比

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

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

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

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

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

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

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

表6：不同訓練方式的分離實(shí)驗

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

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

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