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

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

不同設(shè)備往往對(duì)訓(xùn)練的模型有不同的要求，如模型占用內(nèi)存的大小、推理時(shí)延等。為了滿(mǎn)足設(shè)備的要求，技術(shù)人員往往需要對(duì)大規(guī)模的預(yù)訓(xùn)練模型進(jìn)行壓縮處理。但是由于存在大量復(fù)雜的場(chǎng)景和不同的下游任務(wù)，專(zhuān)門(mén)為每一種場(chǎng)景設(shè)計(jì)壓縮過(guò)的BERT模型，需要投入大量的人力和計(jì)算資源。

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

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

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

表1：之前的 BERT 壓縮工作

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

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

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

搜索空間的架構(gòu)設(shè)計(jì)

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

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

圖1：NAS-BERT 方法框架圖

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

表2：操作集合的設(shè)計(jì)

圖2：卷積操作的設(shè)計(jì)

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

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

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

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

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

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

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

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

實(shí)驗(yàn)結(jié)果

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

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

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

表3：NAS-BERT 和 BERT 對(duì)比

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

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

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

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

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

表5：不同逐漸縮小搜索空間的分離試驗(yàn)

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

表6：不同訓(xùn)練方式的分離實(shí)驗(yàn)

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

表7：部分 NAS-BERT 搜索得到的架構(gòu)

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