可與ViT一較高下，DeepMind從稀疏轉向Soft混合專(zhuān)家模型

對于谷歌 DeepMind 的 Soft MoE，有人表示：「即使它不是萬(wàn)能藥，仍可以算得上一個(gè)突破」。

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

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

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

現在，來(lái)自 Google DeepMind 的研究團隊提出了一種稱(chēng)為「Soft MoE」的新方法 ，解決了許多問(wèn)題。

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

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

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

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

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

此外，該研究用實(shí)驗表明通過(guò)軟路由學(xué)習的表征保留了圖像 - 文本對齊的優(yōu)勢。

Soft MoE 模型

算法描述

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

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

Soft MoE 的特性

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

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

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

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

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

實(shí)現

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

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

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

圖像分類(lèi)實(shí)驗結果

研究者展示了圖像分類(lèi)的三種類(lèi)型的實(shí)驗：

• 訓練帕累托邊界

• 推理時(shí)優(yōu)化模型

• 模型消融

研究者在 JFT-4B 數據集上對模型進(jìn)行預訓練，這是一個(gè)專(zhuān)有數據集，最新版本包含了超過(guò) 4B 張圖像、29k 個(gè)類(lèi)別。

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

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

訓練帕累托 - 優(yōu)化模型

研究者訓練了 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 顯示了每個(gè)類(lèi)別中模型的結果，這些模型位于各自的訓練成本 / 性能帕累托邊界上。在上述兩個(gè)評估指標上，對于任何給定的 FLOP 或時(shí)間預算，Soft MoE 顯著(zhù)優(yōu)于密集和其他稀疏方法。

長(cháng)訓練運行

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

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

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

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

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

從結果來(lái)看，經(jīng)過(guò) 1k TPUv3 days 訓練的 Soft MoE B/16 優(yōu)于在類(lèi)似時(shí)間預算下訓練的 ViT H/14，同時(shí)在 FLOP 推理上成本 10 倍降低，在掛鐘時(shí)間上 5.7 倍減少。

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

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

更多技術(shù)細節和實(shí)驗結果請參閱原論文。