稀疏專(zhuān)家模型是一個(gè)已有 30 年歷史的概念，至今依然被廣泛使用，是深度學(xué)習(xí)中的流行架構(gòu)。此類(lèi)架構(gòu)包括混合專(zhuān)家系統(tǒng)（MoE）、Switch Transformer、路由網(wǎng)絡(luò)、BASE 層等。稀疏專(zhuān)家模型已經(jīng)在自然語(yǔ)言處理、計(jì)算機(jī)視覺(jué)和語(yǔ)音識(shí)別等多個(gè)領(lǐng)域展示出良好的性能。

近日，谷歌 AI 負(fù)責(zé)人 Jeff Dean 等人撰寫(xiě)了一篇稀疏專(zhuān)家模型的綜述，回顧了稀疏專(zhuān)家模型的概念，提供了通用算法的基本描述，最后展望了未來(lái)的研究方向。

論文地址：https://arxiv.org/pdf/2209.01667.pdf

通過(guò)增加計(jì)算預(yù)算、訓(xùn)練數(shù)據(jù)和模型大小，機(jī)器學(xué)習(xí)（尤其是自然語(yǔ)言）取得了顯著進(jìn)步。著名的里程碑語(yǔ)言模型包括 GPT-2（Radford et al., 2018）、BERT（Devlin et al., 2018）、T5（Raffel et al., 2019）、GPT-3（Brown et al., 2020）、Gopher (Rae et al., 2021)、Chinchilla (Hoffmann et al., 2022) 和 PaLM (Chowdhery et al., 2022)。

然而，最先進(jìn)的模型現(xiàn)在需要數(shù)以千計(jì)的專(zhuān)用互連加速器，訓(xùn)練數(shù)周或數(shù)月的時(shí)間，因此這些模型的生產(chǎn)成本很高（Patterson et al., 2021）。隨著機(jī)器學(xué)習(xí)系統(tǒng)規(guī)模的擴(kuò)大，該領(lǐng)域?qū)で蟾行У挠?xùn)練和服務(wù)范式。稀疏專(zhuān)家模型已成為一種很有前途的解決方案。

稀疏專(zhuān)家模型（其中，混合專(zhuān)家系統(tǒng)（MoE）是最流行的變體）是一種特殊的神經(jīng)網(wǎng)絡(luò)，其中一組參數(shù)被劃分為 “專(zhuān)家”，每個(gè)“專(zhuān)家” 具有唯一的權(quán)重。

在訓(xùn)練和推理期間，模型將輸入樣本給到特定的專(zhuān)家權(quán)重，讓每個(gè)樣本僅與網(wǎng)絡(luò)參數(shù)的子集交互，這與將整個(gè)網(wǎng)絡(luò)用于每個(gè)輸入的常規(guī)方法不同。由于每個(gè)樣本僅使用一小部分專(zhuān)家，因此相對(duì)于總模型來(lái)說(shuō)計(jì)算量大幅減小了。

許多現(xiàn)代稀疏專(zhuān)家模型從 Shazeer et al. (2017)中汲取了靈感。該研究訓(xùn)練了當(dāng)時(shí)最大的模型，并取得了最先進(jìn)的語(yǔ)言建模和翻譯結(jié)果。與 Transformer 語(yǔ)言模型結(jié)合使用時(shí)，稀疏專(zhuān)家模型的受歡迎程度進(jìn)一步飆升（Lepikhin et al., 2020; Fedus et al., 2021）。雖然大部分工作都在自然語(yǔ)言處理方面，但稀疏專(zhuān)家模型也已成功用于各種領(lǐng)域，包括計(jì)算機(jī)視覺(jué)（Puigcerver et al., 2020）、語(yǔ)音識(shí)別（You et al., 2021）和多模態(tài)學(xué)習(xí)（Mustafa et al., 2022）。Clark et al. (2022) 研究了稀疏專(zhuān)家模型在不同模型大小和專(zhuān)家數(shù)量下的擴(kuò)展特性。此外，許多基準(zhǔn)測(cè)試的最新結(jié)果目前由 ST-MoE 等稀疏專(zhuān)家模型持有（Zoph et al., 2022）。隨著研究和工程的進(jìn)步，該領(lǐng)域正在迅速發(fā)展。

這篇綜述論文將調(diào)查范圍縮小到狹義深度學(xué)習(xí)時(shí)代（從 2012 年開(kāi)始）的稀疏專(zhuān)家模型，回顧了最近的進(jìn)展并討論了有希望的未來(lái)途徑。

稀疏專(zhuān)家模型

機(jī)器學(xué)習(xí)中的 MoE 概念可以追溯到至少 30 年前，在早期的概念中，專(zhuān)家們定義了一個(gè)完整的神經(jīng)網(wǎng)絡(luò)，MoE 類(lèi)似于一種集成方法。

Eigen et al. (2013) 提出在抖動(dòng)（jittered）的 MNIST 上使用堆疊層專(zhuān)家混合架構(gòu)，這項(xiàng)工作為后來(lái)模型的高效實(shí)現(xiàn)奠定了基礎(chǔ)。

Shazeer et al. (2017)提出在兩個(gè) LSTM 層之間插入一個(gè) MoE 層，由此產(chǎn)生的稀疏模型在機(jī)器翻譯中實(shí)現(xiàn) SOTA 性能。然而，此方法雖然取得了成功，但后續(xù)研究就像冬眠了一樣，停滯不前，大部分研究都轉(zhuǎn)向了 Transformer。

2020 到 2021 年間，GShard 和 Switch Transformer 發(fā)布，兩者都用專(zhuān)家層取代了 Transformer 中的前饋層。

盡管使用一層專(zhuān)家層的方法已經(jīng)成為主導(dǎo)范式，但近兩年的研究重新審視了專(zhuān)家模型作為完全獨(dú)立模型的概念，以實(shí)現(xiàn)模塊化和可組合性。

下圖 2 是 Shazeer et al. (2017) 提出的原始 top-k 路由機(jī)制，這是許多后續(xù)工作的基礎(chǔ)。該綜述論文在第四章中詳細(xì)講解了路由算法的新進(jìn)展。

硬件

現(xiàn)代稀疏專(zhuān)家模型已與用于訓(xùn)練最大神經(jīng)網(wǎng)絡(luò)的分布式系統(tǒng)共同設(shè)計(jì)。

大型神經(jīng)網(wǎng)絡(luò)（Brown et al., 2020; Rae et al., 2021; Chowdhery et al., 2022）的研究表明神經(jīng)網(wǎng)絡(luò)已經(jīng)遠(yuǎn)遠(yuǎn)超過(guò)單個(gè)加速器的內(nèi)存容量，因此權(quán)重，激活函數(shù)，優(yōu)化器變量等張量需要使用各種并行策略分片。

三種常見(jiàn)的并行方法包括：數(shù)據(jù)并行（復(fù)制模型權(quán)重，對(duì)數(shù)據(jù)進(jìn)行分片）、張量模型并行（數(shù)據(jù)和權(quán)重張量跨設(shè)備分片）和 pipeline 并行（整個(gè)層或?qū)咏M被跨設(shè)備分片），混合專(zhuān)家模型通常能夠適應(yīng)這些并行方案。

在MoE 模型的訓(xùn)練和部署方面，Jaszczur et al. (2021)將 Transformer 模型的所有層進(jìn)行稀疏化，從而實(shí)現(xiàn)了 37 倍的推理加速；Kossmann et al. (2022)通過(guò) RECOMPILE 庫(kù)解決了靜態(tài)專(zhuān)家批處理大小的約束問(wèn)題。

除了數(shù)據(jù)并行、模型并行和專(zhuān)家并行外，Rajbhandari et al.(2022) 提出 DeepSpeed-MoE 庫(kù)，以支持 ZeRO partitioning 和 ZeRO-Offload，實(shí)現(xiàn)了 10 倍推理提升和 SOTA 翻譯性能，從而增加了模型在生產(chǎn)服務(wù)中的實(shí)用性。

稀疏專(zhuān)家模型的擴(kuò)展特性

密集神經(jīng)語(yǔ)言模型的交叉熵?fù)p失表現(xiàn)為一種關(guān)于模型參數(shù)計(jì)數(shù)、數(shù)據(jù)量和計(jì)算預(yù)算（Kaplan et al., 2020）的冪律。冪律系數(shù)后來(lái)在 Hoffmann et al. (2022)中得到糾正，這表明計(jì)算最優(yōu)模型需要數(shù)據(jù)和參數(shù)擴(kuò)展更緊密的平衡。相比之下，稀疏專(zhuān)家模型的早期研究以啟發(fā)式方式擴(kuò)展獲得了強(qiáng)有力的實(shí)證結(jié)果，但沒(méi)有仔細(xì)描述擴(kuò)展定律。此外，一些工作強(qiáng)調(diào)了上游（例如預(yù)訓(xùn)練）和下游（例如微調(diào)）行為之間的差異（Fedus et al., 2021; Artetxe et al., 2021），進(jìn)一步復(fù)雜化了對(duì)稀疏專(zhuān)家模型的理解。

上游擴(kuò)展

稀疏專(zhuān)家模型在大型數(shù)據(jù)集上訓(xùn)練時(shí)表現(xiàn)出色。自然語(yǔ)言處理中的常見(jiàn)模式是先執(zhí)行上游訓(xùn)練（例如預(yù)訓(xùn)練），然后對(duì)特定感興趣的數(shù)據(jù)分布進(jìn)行下游訓(xùn)練（例如微調(diào)）。在上游階段，稀疏專(zhuān)家模型一直比密集對(duì)應(yīng)模型產(chǎn)生更高的收益。Shazeer et al. (2017) 在 10 億字語(yǔ)言建?；鶞?zhǔn)（Chelba et al., 2013）上提出了關(guān)于模型參數(shù)和計(jì)算預(yù)算的擴(kuò)展曲線(xiàn)，與密集版本相比取得了顯著收益。Lepikhin et al. (2020)提出了模型擴(kuò)展函數(shù)的改進(jìn)版本，并在其最大的 600B 參數(shù)稀疏模型上獲得了 13.5 的 BLEU 分?jǐn)?shù)增益。Switch Transformer (Fedus et al., 2021) 在 T5 模型上使用相同的計(jì)算資源測(cè)量到 4-7 倍的 wall-time 加速。該工作還研究了作為參數(shù)計(jì)數(shù)函數(shù)的交叉熵?fù)p失擴(kuò)展，但觀察到超過(guò) 256 位專(zhuān)家時(shí)，收益減少了。

下游擴(kuò)展

然而，可靠的上游擴(kuò)展并沒(méi)有立即在下游任務(wù)上產(chǎn)生一致的收益。在一項(xiàng)強(qiáng)調(diào)遷移挑戰(zhàn)的工作中，F(xiàn)edus et al. (2021)使用低計(jì)算、高參數(shù)編碼器 - 解碼器 Transformer（1.6T 參數(shù)，每個(gè)稀疏層有 2048 個(gè)專(zhuān)家）觀察到 4 倍的預(yù)訓(xùn)練改進(jìn)，但它在 SuperGLUE 等密集型推理任務(wù)上的微調(diào)效果不佳。這一發(fā)現(xiàn)暗示了進(jìn)一步的必要研究以及計(jì)算和參數(shù)之間可能需要的平衡。

Du et al.(2021) 展示了稀疏 GLaM 模型的擴(kuò)展，范圍從 1B-64B FLOP，每個(gè)稀疏層使用 64 位專(zhuān)家。GLaM 取得了 SOTA 結(jié)果，在 zero-shot 和 one-shot 性能方面優(yōu)于 175B 參數(shù)的 GPT-3 模型（Brown et al., 2020），同時(shí)在推理時(shí)每個(gè)token使用的 FLOP 減少了 49%，功耗降低了 65% （如下圖4(左)所示）。下圖 4 (右) 是稀疏模型在少樣本推理上表現(xiàn)良好的另一個(gè)例子。

Srivastava et al. (2022) 研究了稀疏模型在多項(xiàng)選擇 BIG-Bench 任務(wù)上的校準(zhǔn)，即度量預(yù)測(cè)概率與正確概率匹配的程度。研究結(jié)果如下圖 5 所示，雖然較大的密集模型和稀疏模型的校準(zhǔn)都得到了改善，但稀疏模型與使用 10 倍以上 FLOP 的密集模型的校準(zhǔn)相當(dāng)。

擴(kuò)展專(zhuān)家層的數(shù)量、大小和頻率

有幾個(gè)重要的超參數(shù)控制著稀疏專(zhuān)家模型的擴(kuò)展，包括：1）專(zhuān)家數(shù)量，2）每個(gè)專(zhuān)家的大小，以及 3）專(zhuān)家層的頻率。這些決策可能對(duì)上游和下游擴(kuò)展產(chǎn)生重大影響。

許多早期的工作擴(kuò)展到每層數(shù)千個(gè)相對(duì)較小的專(zhuān)家，從而產(chǎn)生了出色的預(yù)訓(xùn)練和翻譯質(zhì)量（Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2021）。然而，稀疏模型的質(zhì)量在域偏移（Artetxe et al., 2021）或?qū)Σ煌蝿?wù)分布進(jìn)行微調(diào)（Fedus et al., 2021）時(shí)不成比例地降低。為了使計(jì)算和參數(shù)獲得更好的平衡，用于小樣本推理 (GLaM (Du et al., 2021)) 和微調(diào) (ST-MoE (Zoph et al., 2022)) 的 SOTA 稀疏模型最多僅能使用 64 個(gè)更大的專(zhuān)家。由于專(zhuān)家維度的增加，這些模型需要有針對(duì)加速器的特定系統(tǒng)級(jí)分片策略才能有效運(yùn)行（Du et al., 2021; Rajbhandari et al., 2022）。

路由算法

路由算法是所有稀疏專(zhuān)家架構(gòu)的關(guān)鍵特性，它決定了將樣本發(fā)送到哪里。該領(lǐng)域已被廣泛研究，包括使用固定的、非學(xué)習(xí)路由模式的反直覺(jué)方法（Roller et al., 2021）。由于要對(duì)選擇哪些專(zhuān)家做出離散決策，因此路由決策通常是不可微的。

專(zhuān)家選擇問(wèn)題后來(lái)被重新定義為 Bandit 問(wèn)題，并且已有一些工作使用強(qiáng)化學(xué)習(xí)來(lái)學(xué)習(xí)專(zhuān)家選擇（Bengio et al., 2016; Rosenbaum et al., 2017; 2019; Clark et al., 2022）。Shazeer et al. (2017) 提出了一種可微的啟發(fā)式算法，以避開(kāi)強(qiáng)化學(xué)習(xí)的難點(diǎn)。

該論文詳細(xì)闡述了路由算法的分類(lèi)，并進(jìn)一步說(shuō)明了該領(lǐng)域的關(guān)鍵問(wèn)題——負(fù)載均衡。

稀疏專(zhuān)家模型的迅速發(fā)展

稀疏專(zhuān)家模型的影響正在迅速蔓延到NLP以外的其他領(lǐng)域，包括計(jì)算機(jī)視覺(jué)、語(yǔ)音識(shí)別和多模態(tài)應(yīng)用。盡管領(lǐng)域不同，但模型的架構(gòu)和算法大致相同，下表 1 顯示了不同域的稀疏層輸入表征。

稀疏專(zhuān)家模型近幾年的發(fā)展非常迅速。以NLP領(lǐng)域?yàn)槔琒hazeer et al. (2017) 為 LSTM 語(yǔ)言建模和機(jī)器翻譯引入了混合專(zhuān)家層，這些層被插入到 LSTM 模型的標(biāo)準(zhǔn)層之間。

2020 年，Lepikhin et al. (2020)首次將 MoE 層引入到 Transformer，當(dāng)研究人員將每個(gè)專(zhuān)家層擴(kuò)展到 2048 名專(zhuān)家時(shí)，模型在 100 種不同語(yǔ)言上實(shí)現(xiàn)了 SOTA 翻譯結(jié)果。

Fedus et al. (2021) 創(chuàng)建了一個(gè)稀疏的 1.6T 參數(shù)語(yǔ)言模型，實(shí)現(xiàn)了 SOTA 預(yù)訓(xùn)練質(zhì)量。

新研究正在推動(dòng)小樣本學(xué)習(xí)推理和微調(diào)基準(zhǔn)的發(fā)展。Du et al. (2021) 訓(xùn)練了一個(gè)純 MoE 解碼器語(yǔ)言模型，取得了小樣本上的 SOTA 結(jié)果，并且只需要訓(xùn)練 GPT-3 所需的 1/3 計(jì)算量。Zoph et al. (2022) 提出了 ST-MoE，這是一種稀疏編碼器 - 解碼器模型，在大量推理和生成任務(wù)中都實(shí)現(xiàn)了 SOTA。在 SuperGLUE 上進(jìn)行微調(diào)時(shí)，ST-MoE 的性能優(yōu)于 PaLM-540B ，同時(shí)僅使用約 1/20 的預(yù)訓(xùn)練 FLOP 和 1/40 的推理 FLOP。

什么時(shí)候使用稀疏模型

一個(gè)常見(jiàn)的問(wèn)題是，如果給你一個(gè)固定的計(jì)算或 FLOP 預(yù)算（例如 100 個(gè) GPU 20 小時(shí)），你應(yīng)該訓(xùn)練什么類(lèi)型的模型來(lái)獲得最佳性能？

從根本上說(shuō)，稀疏模型允許通過(guò)增加專(zhuān)家的數(shù)量來(lái)大幅增加模型中的參數(shù)數(shù)量，同時(shí)保持每個(gè)樣本的 FLOP 大致恒定。這種做法有好有壞，具體取決于模型的用途。

當(dāng)你有許多加速器（例如 GPU/TPU）來(lái)承載使用稀疏性帶來(lái)的所有附加參數(shù)時(shí)，稀疏性是很有利的。

使用稀疏性還需要仔細(xì)考慮下游任務(wù)。假如你有許多機(jī)器來(lái)進(jìn)行預(yù)訓(xùn)練，但是用來(lái)微調(diào)或者服務(wù)的機(jī)器卻很少，那么稀疏度（例如專(zhuān)家的數(shù)量）應(yīng)該根據(jù)下游用例中可用的內(nèi)存量進(jìn)行調(diào)整。

在有些情況下，稀疏模型看起來(lái)總是比密集模型差。例如當(dāng)所有參數(shù)都保存在加速器內(nèi)存中，這種情況稀疏模型就不如密集模型了。當(dāng)你有能力在多臺(tái)機(jī)器上并行訓(xùn)練或服務(wù)，以便托管來(lái)自專(zhuān)家的附加模型參數(shù)時(shí)，稀疏模型就非常合適。

此外，該綜述論文還介紹了對(duì)稀疏模型訓(xùn)練的改進(jìn)、可解釋性和未來(lái)的研究方向等，感興趣的小伙伴可以查看原論文，了解更多研究?jī)?nèi)容。