論文出自：Shazeer N, Mirhoseini A, Maziarz K, et al. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer[J]. arXiv preprint arXiv:1701.06538, 2017.
摘要
- 神經網絡的吸收信息的容量（capacity）受限於參數數目。
- 條件計算（conditional computation）針對於每個樣本，激活網絡的部分子網絡進行計算，它在理論上已證明，可以作為一種顯著增加模型容量的方法。
- 在實際中，我們在犧牲少量計算效率的情況下，實現了 1000 倍的模型容量（model capacity）的提升。
- 我們引入了稀疏門控專家混合層（Sparsely-Gated Mixture-of-Experts Layer），包括數以千計的前饋子網絡。對於每一個樣本，有一個可訓練的門控網絡（gating network）會計算這些專家（指前饋子網絡）的稀疏組合。
- 我們把專家混合（MoE）應用於語言建模和機器翻譯任務中，對於這些任務，從訓練語料庫中吸收的巨量知識，是十分關鍵的。
- 在我們提出的模型架構里，MoE 包含 1370 億個參數，以卷積的方式放在堆疊 LSTM 層之間。
- 在大型語言建模和及其翻譯的基准測試中，該模型以更少的計算成本，實現了比最先進方法更好的結果。
1 介紹和相關工作
1.1 條件計算
- 充分利用訓練數據和模型大小的規模，一直以來都是深度學習成功的關鍵。
- 當訓練集足夠大，增加神經網絡的容量（即參數數目），可以得到更高的預測准確度。
- 對於傳統的深度學習模型，對每一個樣本都會激活整個模型，這會導致在訓練成本上，以大約二次方的速度增長，因為模型大小和訓練樣本數目都增加了。
- 當前計算能力和分布式計算的進展，並不能滿足這樣的需求。
- 因此有很多工作提出了各種形式的條件計算，它們在不顯著增加計算成本的情況下，盡量增加模型的容量。
- 在這些算法里，以每個樣本為基礎（on a per-example basis），會激活或凍結網絡中的大部分。
- 這種門控決策機制，可以是二進制的，也可以是稀疏而連續的；可以是隨機性的，也可以是確定性的。
- 門控決策通過有各種形式的強化學習和反向傳播來訓練。
Figure 1：MoE 層嵌入到循環語言模型中。在本例中，稀疏的門控函數選擇兩個專家來執行計算。門控網絡會調整專家的輸出。
- 盡管這種思想在理論上很有前景，但是目前為止，還沒有工作展現在模型容量、訓練時間或模型質量上有足夠的提升。我們把原因歸結為這些挑戰：
- 現代計算設備（特別是 GPU），相比分支（branching）而言，在數值計算上更快。
- 大的批量大小對於性能很關鍵。而條件計算減少了批量大小。
- 網絡帶寬會成為性能瓶頸。
- 損失項可能對於實現好的效果是必需的，因此損失項可能會影響模型質量和負載平衡。
- 對於大型數據集，模型容量是最關鍵的。目前條件計算的文獻處理的圖像識別數據集都相對太小了，難以為大模型提供足夠多的信號。
- 本文首先解決了上述挑戰，並且最后看到了條件計算的前景。
- 我們得到了 1000 倍的模型容量提升，只花費了少量計算開銷
- 得到的結果也優於最頂尖的結果
1.2 本文方法：稀疏門控專家混合層
- 我們的條件計算方法，就是引入了一個新的通用神經網絡組件類型：稀疏門控專家混合層。
- MoE 包含：
- 一些專家，每個專家都是一個簡單的前饋神經網絡。
- 一個可訓練的門控網絡，它會挑選專家的一個稀疏組合，用來處理每個輸入。
- 所有網絡都是使用反向傳播聯合訓練的。
- 盡管該技術是通用的，但是本文聚焦在語言建模和機器翻譯任務中（這些任務都受益於非常大的模型）。
- 具體說來，如圖一所示，我們把 MoE 以卷積的方式（convolutionally）放在多層 LSTM 層之間。
- 在文本的每個位置上，就會調用 MoE 一次，進而可能選擇不同的專家組合。
- 不同的專家會傾向於變得高度專業化（基於語法和語義）。
1.3 專家混合的相關工作
2 混合專家層的結構
MoE 層包括 ：
- $n$ 個“專家網絡”：$E_{1}, \cdots, E_{n}$。
- 一個門控網絡 $G$，其輸出是一個稀疏的 $n$ 維向量。
盡管從理論上講，每個專家網絡只要保持一致的輸入大小和輸出大小就可以了；但是，在本文的研究里，我們限制了專家網絡具有相同的網絡結構，而網絡參數保持獨立。
給定輸入 $x$，定義 $G(x)$ 是門控網絡的輸出；$E_{i}(x)$ 是第 $i$ 個專家網絡的輸出。於是 MoE 模塊的輸出為：
\begin{equation}
y=\sum_{i=1}^{n} G(x)_{i} E_{i}(x)
\end{equation}
基於 $G(x)$ 輸出的稀疏性，我們可以節省計算量。
- 當 $G(x)_{i}=0$ 時，我們無需計算 $E_{i}(x)$。
- 在我們的實驗中，我們有數以千計的專家，但是針對每個樣本，只需要用到少量的專家。
- 如果專家數目非常大，我們可能要采用層次化的 MoE；本文我們不會使用層次化的 MoE，相關細節感興趣可以見附錄 B。
2.1 門控網絡
Softmax Gating
一種朴素的想法是，用一個矩陣乘上輸入，然后經過一個 Softmax 函數，這種方法實際上是一種非稀疏的門控函數：
\begin{equation}
G_{\sigma}(x)=\operatorname{Softmax}\left(x \cdot W_{g}\right)
\end{equation}
Noise Top-K Gating
我們在 Softmax 門控網絡基礎上，加入兩個元素：稀疏性和噪聲。在執行 Softmax 函數之前：
- 我們加入了可調的高斯噪聲
- 噪聲項是為了幫助負載均衡（load balancing），我們在附錄 A 有詳細討論。
- 並且保留前 k 個值
- 這種稀疏性是為了節省計算資源，盡管這種形式的稀疏性，從理論上會造成一些可怕的輸出間斷性，但在實際使用中，我們並沒有觀察到這種問題。
每個分量的噪音量，通過另一個可訓練的權重矩陣 $W_{\text {noise }}$ 來控制。
\begin{equation}
G(x)=\operatorname{Softmax}(\operatorname{KeepTopK}(H(x), k))
\end{equation}
\begin{equation}
H(x)_{i}=\left(x \cdot W_{g}\right)_{i}+\text { StandardNormal }() \cdot \operatorname{Softplus}\left(\left(x \cdot W_{\text {noise }}\right)_{i}\right)
\end{equation}
\begin{equation}
\text { KeepTopK }(v, k)_{i}=\left\{\begin{array}{ll}
v_{i} & \text { if } v_{i} \text { is in the top } k \text { elements of } v . \\
-\infty & \text { otherwise. }
\end{array}\right.
\end{equation}
訓練門控網絡
我們使用簡單的反向傳播來訓練門控網絡以及接下來的模型。
3 解決性能挑戰
3.1 批量減小問題（The Shrinking Batch Problem）
由於門控網絡對每個樣本，在 $n$ 個專家中，選擇 $k$ 個。那么對於 $b$ 個樣本的批次，每個轉接都會收到更加更加小的批次（大概 $\frac{k b}{n} \ll b$）。這會導致朴素的 MoE 實現在專家數量增加時，非常低效。解決批量減小問題，就是需要讓原始的批量大小盡可能的大。然而，批量大小會收到內存的限制。我們提出如下技術來提高批量大小：
- 混合數據並行和模型並行（Mixing Data Parallelism and Model Parallelism）
- 充分利用卷積
- 增加循環 MoE 的批量大小
3.2 網絡帶寬
4 平衡專家的利用率
我們觀察到，門控網絡傾向於收斂到一種不好的狀態，即對相同的少量專家，總是會得到較大的權重。這種不平衡是不斷自我強化的，隨着更好的專家不斷訓練學習，它們更有可能被門控網絡選中。面對這種問題，過去文獻有的用硬性約束，有的用軟性約束。
而我們采用軟性約束方法。我們定義對於一個批次訓練樣本的專家重要度（the importance of an expert），即該專家在一個批次上的門控輸出值的和。並且定義損失項 $L_{\text {importance }}$，加入到模型的總損失上。該損失項等於所有專家重要度的方差的平方，再加上一個手工調節的比例因子 $w_{important}$。這個損失項會鼓勵所有專家有相同的重要度。
\begin{equation}
\text { Importance }(X)=\sum_{x \in X} G(x)
\end{equation}
\begin{equation}
L_{\text {importance }}(X)=w_{\text {importance }} \cdot C V(\text { Importance }(X))^{2}
\end{equation}
盡管現在的損失函數可以保證相同的重要度，專家仍然可能接收到差異很大的樣本數目。例如，某些專家可能接收到少量的大權重的樣本；而某些專家可能接收到更多的小權重的樣本。為了解決這個問題，我們引入了第二個損失函數：$L_{\text {load }}$，它可以保證負載均衡。附錄 A 會包含該函數的定義。
5 實驗
5.1 10 億詞匯的語言建模基准
5.2 1000 億詞匯的谷歌新聞語料庫
5.3 機器翻譯
5.4 多語言機器翻譯
6 結論
- 該工作是第一個展現基於深度網絡的條件計算的重大勝利。
- 我們探討了設計考慮、條件計算的挑戰、從算法和工程上的解決方案。
- 雖然我們聚焦在文本領域上，條件計算仍然可以在其他領域發揮作用。我們期望有更多條件計算的實現和應用。