機器之心報道
編輯：張倩、杜偉
稀疏模型在深度學習領域發揮着越來越重要的作用。對于給定的 token 或樣本，它可以隻激活模型的一小部分，進而在擁有很大的參數量的同時也能做到計算友好。但是，如何可靠地訓練這類模型依然是一個需要解決的問題。在這篇文章中，來自谷歌的 Barret Zoph、Irwan Bello、William Fedus 、Jeff Dean 等研究者給出了一份「高效稀疏專家模型設計指南」。
稀疏專家神經網絡展示了純規模的優勢，并為當今常用的靜态神經網絡架構提供了一種有效的替代方案。稀疏專家網絡不是對所有輸入應用相同的參數，而是為每個輸入動态選擇使用哪些參數。這允許網絡極大地擴充參數的數量，同時保持每個 token 的 FLOPs 大緻不變。這些方法的采用已經帶來了 SOTA 翻譯模型、4-7 倍的預訓練加速，以及僅使用 1/3 的訓練成本就能達到 GPT-3 級的 one-shot 性能。盡管參數數量驚人，但稀疏模型将訓練大型神經網絡的碳足迹降低了一個數量級。然而，困難依然存在。
Fedus et al. (2021)觀察到，與之前的 SOTA 方法（Raffel et al., 2019）相比，稀疏 1.6T 參數模型實作了 4 倍的預訓練加速，但在 SuperGLUE 等常用基準上進行微調時，卻落後于較小的模型。在 Artetxe et al. (2021)中，研究者在域外資料上對 MoE 語言模型進行了微調，并觀察到了相似的差距。
為了解決這一問題，Switch-XXL 模型被提出，該模型參數較少，但計算占用空間增加到原來的 8 倍（FLOPs 大約等于最大的 T5 模型），在自然語言了解任務上的性能有所提高。然而，必要的預訓練受到先前在小規模研究中未檢測到的訓練不穩定性的阻礙。這些不穩定性後來在其他稀疏模型中被識别出來。這些結果揭示了參數和計算的必要平衡，但如何可靠地訓練這種模型依然是一個待解決的問題。
這篇論文的目的就是提高稀疏模型的實用性和可靠性。他們研究了這兩個問題，并給出了設計指南。最後，他們将稀疏模型的參數縮放到 269B，其計算成本與 32B 密集編碼器 - 解碼器 Transformer（穩定、可遷移的 Mixture-of-Experts、ST-MoE-32B）相當。這是稀疏模型首次在遷移學習中實作 SOTA 性能，跨越了一系列不同的任務，包括推理（SuperGLUE、ARC Easy、ARC Challenge)、摘要（XSum、CNN-DM）、閉卷問答（WebQA、Natural Questions）和對抗式構造任務（Winogrande、ANLI R3）。
本文的貢獻可以概括如下：
1、開展了一項關于穩定性技術的品質 - 穩定性權衡（quality-stability trade-offs）大規模研究；
2、引入了 router z-loss 來解決穩定性問題，同時略微提高了模型品質；
3、給出了關于稀疏和密集模型的微調分析，揭示了二者對批大小和學習率的不同超參數敏感性；他們發現，糟糕的超參數導緻密集模型上幾乎沒有微調增益，盡管預訓練有很大的加速；
4、給出了分布式環境下設計 Pareto 高效稀疏模型的架構、routing 和模型設計原則；
5、給出了追蹤跨專家層的 token routing 決策的定性分析；
6、訓練出了一個 269B 稀疏模型，在一組不同的自然語言基準上實作了 SOTA 性能。
router z-loss
穩定神經網絡最成功的方法之一是對激活的限制和梯度。一種流行的方法是在通過深度網絡反向傳播時，裁剪梯度範數來彌補爆炸梯度。
在這篇論文中，研究者使用 Adafactor 優化器是因為它的記憶體效率（盡管最近推出的 8 位優化器 (Dettmers et al., 2021) 可能會提供更好的 trade-off）。Adafactor 使用更新裁剪（update clipping），而不是梯度裁剪（gradient clipping），其中對權重的更改被限制在一定的範數以下。他們嘗試将更新裁剪收緊到更小的值。
接下來，他們研究了即将進入 router 的 logit 上的限制。router 以 float32 計算專家的機率分布。然而，研究者發現，在最大的規模下，這不足以帶來可靠的訓練結果。為了解決這個問題，他們引入了 router z-loss，
其中，B 是 token 的數目，N 是專家數，x ∈ RB×N 是将要進入 router 的 logit。
下表 4 顯示，在三次運作中，update clipping 和 router z-loss 都穩定了模型，但是 update clipping 嚴重影響了模型的品質。是以，研究者使用 z-loss 方法來固定模型穩定性。
router z-loss 引入了另一個超參數 (c_z)，這是一個權重系數，作為優化的總損失的一部分。總損失是交叉熵損失 (cross entropy loss, L_CE)、輔助負載平衡損失 (auxiliary load balance loss, L_B) 和 router z-loss (L_Z) 的線性權重組合。
基于用超參數掃描進行預訓練後的最佳模型品質，研究者選擇 c_z = 0.001 的值。附錄 B 記錄了預訓練過程中的損失。
稀疏模型的設計
密集模型的設計受到 Kaplan et al. (2020) 的基礎工作的指導。但是到了稀疏模型這裡，還有無數的額外問題需要解決，比如：(1)使用多少專家？(2)使用哪種 routing 算法？(3)容量因子（capacity factor）的值是多少？(4)硬體如何改變這些決策？在本文中，研究者給出的建議是：
1、在他們的設定中，他們推薦容量因子為 1.25 的 top-2 routing，每個核心最多有一個專家；
2、在評估期間，可以更改容量因子，以适應新的記憶體 / 計算要求；
3、密集層疊加和乘法偏置（multiplicative bias）可以提高品質。
更多細節請參考原論文。
論文連結：https://arxiv.org/pdf/2202.08906.pdf