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

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

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

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

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

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

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

表1：之前的 BERT 壓縮工作

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

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

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

搜索空間的架構設計

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

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

圖1：NAS-BERT 方法框架圖

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

表2：操作集合的設計

圖2：卷積操作的設計

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