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

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

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

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

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

現(xiàn)在，來自 Google DeepMind 的研究團(tuán)隊(duì)提出了一種稱為「Soft MoE」的新方法 ，解決了許多問題。

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

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

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

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

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

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

Soft MoE 模型

算法描述

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

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

Soft MoE 的特性

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

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

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

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

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

實(shí)現(xiàn)

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

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

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

圖像分類實(shí)驗(yàn)結(jié)果

研究者展示了圖像分類的三種類型的實(shí)驗(yàn)：

• 訓(xùn)練帕累托邊界

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

• 模型消融

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

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

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

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

研究者訓(xùn)練了 VIT-S/8、VIT-S/16、VIT-S/32、VIT-B/16、VIT-B/32、VIT-L/16、VIT-L/32 和 VIT-H/14 模型，以及它們的稀疏對(duì)應(yīng)模型。

下圖 3a 和 3b 顯示了每個(gè)類別中模型的結(jié)果，這些模型位于各自的訓(xùn)練成本 / 性能帕累托邊界上。在上述兩個(gè)評(píng)估指標(biāo)上，對(duì)于任何給定的 FLOP 或時(shí)間預(yù)算，Soft MoE 顯著優(yōu)于密集和其他稀疏方法。

長(zhǎng)訓(xùn)練運(yùn)行

除了較短的運(yùn)行和消融實(shí)驗(yàn)之外，研究者還訓(xùn)練了運(yùn)行更長(zhǎng)（幾百萬 step）的模型，從而在更大的計(jì)算規(guī)模上測(cè)試 Soft MoE 的性能。

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

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

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

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

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

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

此外，Soft MoE L/16 大幅地?fù)魯∷心Ｐ?，同時(shí)推理速度是 ViT H/14 的近 2 倍。

更多技術(shù)細(xì)節(jié)和實(shí)驗(yàn)結(jié)果請(qǐng)參閱原論文。