隨著大模型涌現出令人驚艷的性能，模型大小已經成為影響模型性能的關鍵因素之一。通常，對 Transformer 模型來說，模型越大，性能就會越好，但計算成本也會增加。近期有研究表明，模型大小和訓練數據必須一起擴展，才能最佳地使用給定的訓練計算預算。

稀疏混合專家模型（MoE）是一種很有前途的替代方案，可以在計算成本較少的情況下，擴展模型的大小。稀疏 MoE Transformer 有一個關鍵的離散優化問題：決定每個輸入 token 應該使用哪些模塊。這些模塊通常是稱為專家的 MLP。

為了讓 token 與專家良好匹配，人們設計了許多方法，包括線性規劃、強化學習、最優傳輸（optimal transport）等。在許多情況下，需要啟發式輔助損失（auxiliary loss）來平衡專家的利用率并最大限度地減少未分配的 token。這些挑戰在分布外場景中可能會加劇。

現在，來自 Google DeepMind 的研究團隊提出了一種稱為「Soft MoE」的新方法 ，解決了許多問題。

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

Soft MoE 不采用稀疏且離散的路由器在 token 和專家之間進行硬分配（hard assignment），而是通過混合 token 來執行軟分配（soft assignment）。值得注意的是，這種方法會計算所有 token 的多個加權平均值（weighted average），其中權重取決于 token 和專家，然后由相應的專家處理每個加權平均值。

常見的稀疏 MoE 算法通常會學習一些路由器參數，但這些算法的效果有時甚至不如隨機固定路由。在 Soft MoE 中，由于每個路由（或混合）參數都是根據單個輸入 token 直接更新的，因此可以在訓練路由器期間提供穩定性。研究團隊還觀察到，在訓練期間，大部分輸入 token 可以同時改變網絡中的離散路由。

此外，硬路由（hard routing）在專家模塊數量較多時可能具有挑戰性，因此大多數研究的訓練只有幾十個專家模塊。相比之下，Soft MoE 可擴展至數千個專家模塊，并且可以通過構建實現平衡。最后，Soft MoE 在推理時不存在批次效應（batch-effect）。

該研究進行了一系列實驗來探究 Soft MoE 方法的實際效果。實驗結果表明，Soft MoE L/16 在上游任務、少樣本任務和微調方面擊敗了 ViT H/14，并且 Soft MoE L/16 僅需要一半的訓練時間，推理速度還是 ViT H/14 的 2 倍。值得注意的是，盡管 Soft MoE B/16 的參數量是 ViT H/14 的 5.5 倍，但 Soft MoE B/16 的推理速度卻是 ViT H/14 的 5.7 倍。

此外，該研究用實驗表明通過軟路由學習的表征保留了圖像 - 文本對齊的優勢。

Soft MoE 模型

算法描述

Soft MoE 路由算法如下圖 2 所示。研究者使用 X ∈ R^m×d 來表示一個序列的輸入 token，其中 m 是 token 數量，d 是維數。每個 MoE 層使用一組 n 個專家函數應用于單個 token，即 {f_i : R^d → R^d}_1:n。每個專家將處理 p 個 slot，每個 slot 有相應的 d 維參數向量。他們用來表示這些參數。

研究者遵循稀疏 MoE 的常規設計，利用 Soft MoE 塊替換掉了 Transformer MLP 塊的一個子集。這里通常會替換掉 MLP 塊的后半部分。slot 的總數量是 Soft MoE 層的關鍵超參數，這是因為時間復雜度依賴于 slot 數量而不是專家數量。比如可以將 slot 數設置為與輸入序列長度相等，以匹配等效密集 Transformer 的 FLOP。

Soft MoE 的特性

首先 Soft MoE 完全可微。Soft MoE 中的所有操作都是連續且完全可微的。我們可以將帶有 softmax 分數的加權平均值解釋為軟分類，這也是 Soft MoE 算法名稱的由來。作為對比，稀疏 MoE 方法通常采用的是硬分類。

其次 Soft MoE 沒有 token dropping 和專家不平衡。Soft MoE 基本上不受這兩點的影響，這得益于每個 slot 都填充了所有 token 的加權平均值。并且由于 softmax，所有權重都是嚴格正的。

再次 Soft MoE 速度快。它的主要優勢是完全避免了排序或 top-k 操作，這些操作速度慢并且通常不太適合硬件加速器。因此，Soft MoE 的速度明顯要快于大多數稀疏 MoE，具體如下圖 6 所示。

Soft MoE 還兼具稀疏和密集的特點。稀疏 MoE 的稀疏性來自于專家參數僅應用于輸入 token 的子集。然而，Soft MoE 在技術上不稀疏，這是因為每個 slot 是所有輸入 token 的加權平均值。并且每個輸入 token 會極少部分激活所有模型參數。同樣所有輸出 token 也極少部分依賴所有 slot（和專家）。還要注意一點，Soft MoE 不是密集 MoE（其中每個專家處理所有輸入 token），它的每個專家僅處理 slot 的子集。

最后 Soft MoE 具有序列性。由于它組合了每個輸入序列中的所有 token，因此只需將組大小設置為一個大序列。每個專家會處理每個輸入的 token，這可能會在一定程度上限制高級 specialization 的數量。這也意味著 Soft MoE 呈現逐實例確定性和速度快，而稀疏 MoE 的典型實例不是這樣。

實現

時間復雜度。假設單個專家函數的逐 token 成本是 O (k)，則一個 Soft MoE 層的時間復雜度為 O (mnpd + npk)。通過為每個專家選擇 p = O (m/n) 個 slot，也就是 token 數量除以專家數量，成本可以降低至 O (m^2d + mk)。

歸一化。在 Transformer 中，MoE 層通常用來替換掉每個編碼器塊中的前饋層。因此當使用預歸一化作為大多數現代 Transformer 架構時，MoE 層的輸入是「層歸一化的」。

分布式模型。研究者采用標準技術將模型分布在很多設備上。分布式模型通常會增加模型的成本開銷，不過他們上文推導的基于 FLOP 的時間復雜度分析并沒有捕獲這一點。因此在所有實驗中，研究者不僅測量了 FLOP，還測量了以 TPUv3-chip-hour 為單位的掛鐘時間。

圖像分類實驗結果

研究者展示了圖像分類的三種類型的實驗：

• 訓練帕累托邊界
• 推理時優化模型
• 模型消融

研究者在 JFT-4B 數據集上對模型進行預訓練，這是一個專有數據集，最新版本包含了超過 4B 張圖像、29k 個類別。

在預訓練期間，他們提供了兩個指標的評估結果，即 JFT-4B 的上游驗證 precision-at-1 和 ImageNet 10-shot 準確率。此外還提供了在 ImageNet-1k（1.3M 張圖像）上進行微調后，ImageNet-1k 驗證集上的準確率。

研究者對比了兩個流行 MoE 路由算法，分別是 Tokens Choice 和 Experts Choice。

訓練帕累托 - 優化模型

研究者訓練了 VIT-S/8、VIT-S/16、VIT-S/32、VIT-B/16、VIT-B/32、VIT-L/16、VIT-L/32 和 VIT-H/14 模型，以及它們的稀疏對應模型。

下圖 3a 和 3b 顯示了每個類別中模型的結果，這些模型位于各自的訓練成本 / 性能帕累托邊界上。在上述兩個評估指標上，對于任何給定的 FLOP 或時間預算，Soft MoE 顯著優于密集和其他稀疏方法。

長訓練運行

除了較短的運行和消融實驗之外，研究者還訓練了運行更長（幾百萬 step）的模型，從而在更大的計算規模上測試 Soft MoE 的性能。

首先研究者訓練了從 Small 到 Huge 的不同大小的 ViT 和 Soft MoE 模型，它們運行了 4 百萬 step。下圖 4 和表 2 展示了結果。

其中圖 4 展示了 Soft MoE 與 ViT 的 JFT-4B 精度、ImageNet 10-shot 準確率和 ImageNet 微調準確率，以及 ExaFLOPS 的訓練成本。表 2 提供了所有結果。對于給定的計算預算，Soft MoE 模型的性能遠優于 ViT 模型。

再來看針對推理進行優化的 Soft MoE。對于具有較小主干的 Soft MoE 可以與較大 ViT 模型一較高下這一事實，研究者受到了鼓舞，繼續訓練小的主干，以在非常低的推理成本下獲得更高質量的模型。

對于更長時間的運行，研究者觀察到冷卻時間（學習率線性降低到 0）越長，Soft MoE 的效果很好。因此，他們將冷卻時長從 50k step 增加到最多 500k。下圖 5 展示了這些模型。

從結果來看，經過 1k TPUv3 days 訓練的 Soft MoE B/16 優于在類似時間預算下訓練的 ViT H/14，同時在 FLOP 推理上成本 10 倍降低，在掛鐘時間上 5.7 倍減少。

即使將 ViT-H/14 的訓練預算加倍（2M step 和 2039.8 train days），Soft MoE B/16（1011.4 days）也幾乎與它性能相近。

此外，Soft MoE L/16 大幅地擊敗所有模型，同時推理速度是 ViT H/14 的近 2 倍。

更多技術細節和實驗結果請參閱原論文。