Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer; Azalia Mirhoseini; Krzysztof Maziarz; Andy Davis; Quoc Le; Geoffrey Hinton; Jeff Dean

論文 · arXiv 1701.06538

極大規模神經網路：稀疏門控專家混合層 (Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer)

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean

發表於 2017-01-23cs.LGcs.CLcs.NEstat.ML arXiv PDF

用 XYZNOT 閱讀（全文翻譯與 AI 講解）

三句話總結

🎯 密集神經網路在擴展模型容量以吸收海量資料集時，面臨運算成本二次方暴增的問題。
🔧 引入稀疏門控專家混合（MoE）層，在分散式 GPU 分片間採用 Noisy Top-K 門控與輔助負載平衡損失。
📈 將語言模型與翻譯模型擴展至 1370 億個參數，以更低的運算成本將最優已發表 LM 困惑度降低了 24% 至 39%。

摘要

神經網路吸收資訊的能力受限於其參數數量。條件運算（Conditional computation）在理論上被提出作為大幅增加模型容量而不成比例增加運算量的方法，即針對每個樣本僅啟動網路的一部分。然而在實踐中，存在顯著的演算法與效能挑戰。在本研究中，我們解決了這些挑戰，終於實現了條件運算的承諾，在現代 GPU 叢集上僅損失極少運算效率的情況下，實現了超過 1000 倍的模型容量提升。我們提出了一種稀疏門控專家混合層（Sparsely-Gated Mixture-of-Experts layer, MoE），由多達數千個前饋子網路組成。可訓練的門控網路為每個樣本決定要使用的稀疏專家組合。我們將 MoE 應用於語言模型與機器翻譯任務，在這些任務中，模型容量對於吸收訓練語料庫中大量的知識至關重要。我們提出的模型架構中，擁有高達 1370 億參數的 MoE 以卷積方式應用於堆疊的 LSTM 層之間。在大型語言模型與機器翻譯基準測試中，這些模型以更低的運算成本取得了顯著優於現有最先進技術（SOTA）的成果。

英文原文摘要

The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation. In practice, however, there are significant algorithmic and performance challenges. In this work, we address these challenges and finally realize the promise of conditional computation, achieving greater than 1000x improvements in model capacity with only minor losses in computational efficiency on modern GPU clusters. We introduce a Sparsely-Gated Mixture-of-Experts layer (MoE), consisting of up to thousands of feed-forward sub-networks. A trainable gating network determines a sparse combination of these experts to use for each example. We apply the MoE to the tasks of language modeling and machine translation, where model capacity is critical for absorbing the vast quantities of knowledge available in the training corpora. We present model architectures in which a MoE with up to 137 billion parameters is applied convolutionally between stacked LSTM layers. On large language modeling and machine translation benchmarks, these models achieve significantly better results than state-of-the-art at lower computational cost.

論文解讀

問題與動機

深度學習模型在大型資料集上訓練時，能極大地受益於參數容量。然而，標準架構會為每個輸入 token 或樣本啟動整個網路。隨著網路規模和資料集體量的擴大，這會導致運算成本出現不切實際的增長。

雖然條件運算（每個樣本僅啟動模型的子集）在理論上已被提出，但早期的實作在現代加速器上面臨以下困境：

GPU 更偏好均勻的算術運算，而非分支跳轉。
「縮小批次問題」（Shrinking batch problem）：將樣本子集路由後，會導致單個子網路面臨極小且低效的批次大小。
裝置間的網路頻寬瓶頸。
路由不平衡：少數子網路佔據主導地位，而其他子網路則處於飢餓狀態。

核心概念

稀疏門控 MoE 層（Sparsely-Gated MoE Layer）： 前饋子網路（「專家」）被放置在循環層（如堆疊的 LSTM）之間。每個位置僅評估稀疏的 top- $k$ 專家選擇。
Noisy Top-K 門控（Noisy Top-K Gating）： 在進行 top- $k$ 選擇和 softmax 之前，將可調的高斯雜訊加入門控 logits 中，以提供連續且可反向傳播的梯度，並促進專家之間的均衡探索。
結合資料與模型平行： 標準層和門控在不同裝置間以資料平行方式運作，而專家則以模型平行方式劃分在各個 GPU 上。所有裝置的批次會被合併，使專家能夠接收足夠大的聚合 mini-batches。
平衡約束： 針對專家重要性與專家負載施加明確的輔助損失，防止門控網路塌縮至過度使用少數專家的退化解。

方法（與關鍵公式）

MoE 層由 $n$ 個前饋專家網路 $E_{1}, \dots, E_{n}$ 和一個門控網路 $G (x)$ 組成。該層的輸出為： $y = i = 1 \sum n G (x)_{i} E_{i} (x)$ 只要 $G (x)_{i} = 0$ ，就可以避免相應的運算。

Noisy Top-K 門控

為了將稀疏性與可微門控相結合，雜訊被注入到線性投影中： $H (x)_{i} = (x \cdot W_{g})_{i} + StandardNormal () \cdot Softplus ((x \cdot W_{noise})_{i})$ $KeepTopK (v, k)_{i} = {v_{i} - \infty 若 v_{i} 屬於 v 的前 k 個元素其他情況$ $G (x) = Softmax (KeepTopK (H (x), k))$

負載平衡損失

為了強制執行平等的總門控權重，使用變異係數（ $C V$ ）在批次 $X$ 上定義重要性損失： $Importance (X) = x \in X \sum G (x)$ $L_{importance} (X) = w_{importance} \cdot C V (Importance (X))^{2}$ 為了確保每個專家的樣本數量均勻，平滑負載估計器使用常態累積分布函數（CDF） $Φ$ ： $P (x, i) = Φ (\frac{( x \cdot W _{g} ) _{i} - kth_excluding ( H ( x ) , k , i )}{Softplus (( x \cdot W _{noise} ) _{i} )})$ $Load (X)_{i} = x \in X \sum P (x, i)$ $L_{load} (X) = w_{load} \cdot C V (Load (X))^{2}$

結果（關鍵數據）

10 億詞基準（1 Billion Word Benchmark）： 在運算量相當的情況下，MoE-4096-h 模型的測試困惑度達到 34.1（與運算量相當的基準 LSTM 相比降低了 24%）。高運算預算的 MoE 模型在 32 塊 K40 GPU 上僅耗時 47 小時即達到 28.0 的測試困惑度，優於先前耗時 59 小時達到 34.7 的 SOTA 結果。
1000 億詞語料庫： 擴展至 65,536 個專家（680 億參數）時，測試困惑度較密集基準降低了 39%（28.9 對比 47.0），並維持了 0.72 TFLOPS/GPU 的運算效率。
機器翻譯： 經 MoE 增強的 GNMT 在 WMT'14 En $\to$ Fr 上達到 40.56 BLEU（比無強化學習的 GNMT 高出 +1.34 BLEU），在 WMT'14 En $\to$ De 上達到 26.03 BLEU（+1.12 BLEU），且每個時間步所需的運算量不到一半。
多語言翻譯： 單個多語言 MoE 模型在 12 個語言對中的 11 個上均優於多語言 GNMT 基準（在韓語 $\to$ 英語上提升高達 +5.84 BLEU），並將驗證集困惑度降低了 19%。

限制與開放性問題

極端稀疏性可能會導致報酬遞減或效能衰退；在 100B 詞資料集上，當專家數從 65,536 擴展到 131,072 時，困惑度反而變差。
分散式通訊開銷仍高度依賴於隱藏維度與輸入/輸出維度的比例，導致較小容量的專家在頻寬利用上效率較低。
在 RNN 循環步驟內部重複應用 MoE 會破壞跨序列位置的批次展開（batch-unrolling）最佳化，需要採用活化值重算機制。

關鍵術語

Sparsely-Gated Mixture-of-Experts Layer (MoE)（稀疏門控專家混合層）: 一種神經網路層，包含多個平行專家子網路，並由可訓練的門控網路為每個輸入樣本僅選擇稀疏的專家子集。
Conditional Computation（條件運算）: 一種架構範式，其中神經網路的某些部分會根據每個樣本動態啟動或停用，從而增加參數容量而不會成比例增加運算量。
Noisy Top-K Gating（含雜訊 Top-K 門控）: 一種門控機制，在僅保留前 k 個值並套用 softmax 之前，向 logits 添加可調的高斯雜訊，從而實現反向傳播並防止路由塌縮。
Shrinking Batch Problem（批次縮小問題）: 當輸入被稀疏分配到眾多專家時，個別子網路接收到的有效批次大小會縮小，從而降低硬體的算術運算效率。
Importance Loss（重要性損失）: 一種輔助損失，懲罰整個批次中每個專家的門控值總和之變異係數平方，以促使專家分配均勻。
Load Loss（負載損失）: 一種輔助目標，透過最小化專家間估計機率加權 token 分配數量的變異數，以平衡硬體負載。
Hierarchical Mixture of Experts（階層式專家混合）: 一種兩級 MoE 結構，其中一級門控網路將輸入路由到二級 MoE 叢集，以減少分支因子並降低分散式通訊開銷。

常見問題

門控網路如何允許離散的 top-k 路由進行反向傳播？

因為 top-k 選擇是在 softmax 之前進行的，被選中專家的非零門控值相對於輸入權重仍保留連續導數，從而允許梯度反向傳播通過門控網路。

如何在分散式裝置間解決批次縮小問題？

該模型混合了資料平行與模型平行：標準層在各裝置間處理獨立的 mini-batches，但所有裝置的 token 會在路由至託管於不同 GPU 上的固定專家之前被聚合在一起。

為什麼門控網路中需要雜訊？

注入的高斯雜訊可以在訓練早期促進對各專家的探索，並為負載平衡損失提供平滑的估計器，防止自我強化的專家塌縮現象。

當模型擴展到 131,072 個專家時會發生什麼事？

困惑度從 28.9（65,536 個專家）略微退化至 29.2（131,072 個專家），運算效率降至 0.30 TFLOPS/GPU，這可能是由於極端稀疏性以及固定且未相應放大的批次大小所致。

本頁解讀由 AI 根據論文生成，可能有錯誤；細節請以原文為準。

用 XYZNOT 閱讀（全文翻譯與 AI 講解）