稀疏專家模型是一個已有 30 年歷史的概念，至今依然被廣泛使用，是深度學習中的流行架構。此類架構包括混合專家系統（MoE）、Switch Transformer、路由網絡、BASE 層等。稀疏專家模型已經在自然語言處理、計算機視覺和語音識別等多個領域展示出良好的性能。

近日，谷歌 AI 負責人 Jeff Dean 等人撰寫了一篇稀疏專家模型的綜述，回顧了稀疏專家模型的概念，提供了通用算法的基本描述，最后展望了未來的研究方向。

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

通過增加計算預算、訓練數據和模型大小，機器學習（尤其是自然語言）取得了顯著進步。著名的里程碑語言模型包括 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)。

然而，最先進的模型現在需要數以千計的專用互連加速器，訓練數周或數月的時間，因此這些模型的生產成本很高（Patterson et al., 2021）。隨著機器學習系統規模的擴大，該領域尋求更有效的訓練和服務范式。稀疏專家模型已成為一種很有前途的解決方案。

稀疏專家模型（其中，混合專家系統（MoE）是最流行的變體）是一種特殊的神經網絡，其中一組參數被劃分為 “專家”，每個“專家” 具有唯一的權重。

在訓練和推理期間，模型將輸入樣本給到特定的專家權重，讓每個樣本僅與網絡參數的子集交互，這與將整個網絡用于每個輸入的常規方法不同。由于每個樣本僅使用一小部分專家，因此相對于總模型來說計算量大幅減小了。

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

這篇綜述論文將調查范圍縮小到狹義深度學習時代（從 2012 年開始）的稀疏專家模型，回顧了最近的進展并討論了有希望的未來途徑。

稀疏專家模型

機器學習中的 MoE 概念可以追溯到至少 30 年前，在早期的概念中，專家們定義了一個完整的神經網絡，MoE 類似于一種集成方法。

Eigen et al. (2013) 提出在抖動（jittered）的 MNIST 上使用堆疊層專家混合架構，這項工作為后來模型的高效實現奠定了基礎。

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

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

盡管使用一層專家層的方法已經成為主導范式，但近兩年的研究重新審視了專家模型作為完全獨立模型的概念，以實現模塊化和可組合性。

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

硬件

現代稀疏專家模型已與用于訓練最大神經網絡的分布式系統共同設計。

大型神經網絡（Brown et al., 2020; Rae et al., 2021; Chowdhery et al., 2022）的研究表明神經網絡已經遠遠超過單個加速器的內存容量，因此權重，激活函數，優化器變量等張量需要使用各種并行策略分片。

三種常見的并行方法包括：數據并行（復制模型權重，對數據進行分片）、張量模型并行（數據和權重張量跨設備分片）和 pipeline 并行（整個層或層組被跨設備分片），混合專家模型通常能夠適應這些并行方案。

在MoE 模型的訓練和部署方面，Jaszczur et al. (2021)將 Transformer 模型的所有層進行稀疏化，從而實現了 37 倍的推理加速；Kossmann et al. (2022)通過 RECOMPILE 庫解決了靜態專家批處理大小的約束問題。

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

稀疏專家模型的擴展特性

密集神經語言模型的交叉熵損失表現為一種關于模型參數計數、數據量和計算預算（Kaplan et al., 2020）的冪律。冪律系數后來在 Hoffmann et al. (2022)中得到糾正，這表明計算最優模型需要數據和參數擴展更緊密的平衡。相比之下，稀疏專家模型的早期研究以啟發式方式擴展獲得了強有力的實證結果，但沒有仔細描述擴展定律。此外，一些工作強調了上游（例如預訓練）和下游（例如微調）行為之間的差異（Fedus et al., 2021; Artetxe et al., 2021），進一步復雜化了對稀疏專家模型的理解。

上游擴展

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

下游擴展

然而，可靠的上游擴展并沒有立即在下游任務上產生一致的收益。在一項強調遷移挑戰的工作中，Fedus et al. (2021)使用低計算、高參數編碼器 - 解碼器 Transformer（1.6T 參數，每個稀疏層有 2048 個專家）觀察到 4 倍的預訓練改進，但它在 SuperGLUE 等密集型推理任務上的微調效果不佳。這一發現暗示了進一步的必要研究以及計算和參數之間可能需要的平衡。

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

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

擴展專家層的數量、大小和頻率

有幾個重要的超參數控制著稀疏專家模型的擴展，包括：1）專家數量，2）每個專家的大小，以及 3）專家層的頻率。這些決策可能對上游和下游擴展產生重大影響。

許多早期的工作擴展到每層數千個相對較小的專家，從而產生了出色的預訓練和翻譯質量（Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2021）。然而，稀疏模型的質量在域偏移（Artetxe et al., 2021）或對不同任務分布進行微調（Fedus et al., 2021）時不成比例地降低。為了使計算和參數獲得更好的平衡，用于小樣本推理 (GLaM (Du et al., 2021)) 和微調 (ST-MoE (Zoph et al., 2022)) 的 SOTA 稀疏模型最多僅能使用 64 個更大的專家。由于專家維度的增加，這些模型需要有針對加速器的特定系統級分片策略才能有效運行（Du et al., 2021; Rajbhandari et al., 2022）。

路由算法

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

專家選擇問題后來被重新定義為 Bandit 問題，并且已有一些工作使用強化學習來學習專家選擇（Bengio et al., 2016; Rosenbaum et al., 2017; 2019; Clark et al., 2022）。Shazeer et al. (2017) 提出了一種可微的啟發式算法，以避開強化學習的難點。

該論文詳細闡述了路由算法的分類，并進一步說明了該領域的關鍵問題——負載均衡。

稀疏專家模型的迅速發展

稀疏專家模型的影響正在迅速蔓延到NLP以外的其他領域，包括計算機視覺、語音識別和多模態應用。盡管領域不同，但模型的架構和算法大致相同，下表 1 顯示了不同域的稀疏層輸入表征。

稀疏專家模型近幾年的發展非常迅速。以NLP領域為例，Shazeer et al. (2017) 為 LSTM 語言建模和機器翻譯引入了混合專家層，這些層被插入到 LSTM 模型的標準層之間。

2020 年，Lepikhin et al. (2020)首次將 MoE 層引入到 Transformer，當研究人員將每個專家層擴展到 2048 名專家時，模型在 100 種不同語言上實現了 SOTA 翻譯結果。

Fedus et al. (2021) 創建了一個稀疏的 1.6T 參數語言模型，實現了 SOTA 預訓練質量。

新研究正在推動小樣本學習推理和微調基準的發展。Du et al. (2021) 訓練了一個純 MoE 解碼器語言模型，取得了小樣本上的 SOTA 結果，并且只需要訓練 GPT-3 所需的 1/3 計算量。Zoph et al. (2022) 提出了 ST-MoE，這是一種稀疏編碼器 - 解碼器模型，在大量推理和生成任務中都實現了 SOTA。在 SuperGLUE 上進行微調時，ST-MoE 的性能優于 PaLM-540B ，同時僅使用約 1/20 的預訓練 FLOP 和 1/40 的推理 FLOP。

什么時候使用稀疏模型

一個常見的問題是，如果給你一個固定的計算或 FLOP 預算（例如 100 個 GPU 20 小時），你應該訓練什么類型的模型來獲得最佳性能？

從根本上說，稀疏模型允許通過增加專家的數量來大幅增加模型中的參數數量，同時保持每個樣本的 FLOP 大致恒定。這種做法有好有壞，具體取決于模型的用途。

當你有許多加速器（例如 GPU/TPU）來承載使用稀疏性帶來的所有附加參數時，稀疏性是很有利的。

使用稀疏性還需要仔細考慮下游任務。假如你有許多機器來進行預訓練，但是用來微調或者服務的機器卻很少，那么稀疏度（例如專家的數量）應該根據下游用例中可用的內存量進行調整。

在有些情況下，稀疏模型看起來總是比密集模型差。例如當所有參數都保存在加速器內存中，這種情況稀疏模型就不如密集模型了。當你有能力在多臺機器上并行訓練或服務，以便托管來自專家的附加模型參數時，稀疏模型就非常合適。

此外，該綜述論文還介紹了對稀疏模型訓練的改進、可解釋性和未來的研究方向等，感興趣的小伙伴可以查看原論文，了解更多研究內容。