Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

3392 words

17 minutes

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

2025-07-05

Research

Transformer

/

MoE

/

sparse

/

Paper

一、前言#

論文中文-《切換變換器：通過簡單高效的稀疏性機制擴展至萬億參數模型》，請搭配論文原文食用。

因為目前我需要針對黑洞問題去設計一個 MoE的架構。

二、研究背景與動機 (Background and Motivation)#

在 Switch Transformer 誕生之前 (2021年初)，深度學習領域正處於一個「暴力美學」的時代。以 GPT-3 為代表的 密集型 (Dense) 模型 證明了一個簡單粗暴的真理：模型的參數越多、數據越多、計算資源越多，模型就越強大。研究人員發現，模型性能與其參數數量、數據集大小和計算預算之間存在著冪次法則 (Power-law) 關係。

然而，這種「大力出奇蹟」的模式也帶來了巨大的問題：

驚人的計算成本： 訓練一個數千億參數的密集模型，需要消耗海量的運算能力 (FLOPS) 和電力，成本極其高昂。
效率瓶頸： 在密集模型中，對於任何輸入，模型的所有參數都需要參與計算。這就像無論問題難易，都要召集公司全員開會，效率極低。

學術界一直在尋找更高效的擴展方式。其中一個古老而優雅的思想是 專家混合模型 (Mixture of Experts, MoE)。這個想法最早在 1990 年代就被提出，其核心理念是：與其讓一個龐大的「通才」模型解決所有問題，不如訓練一大群各有所長的「專家」模型，每次只根據輸入的特性，選擇性地激活一小部分最相關的專家來處理。

這樣做的好處是顯而易見的：模型的總參數可以非常龐大（專家數量可以很多），但處理每個輸入的實際計算成本卻是恆定的，因為只有少數專家在工作。

儘管 MoE 的概念很吸引人，但在 Switch Transformer 之前，它的廣泛應用一直受阻，主要因為三大障礙：(1) 複雜性高、(2) 跨設備通訊成本巨大、(3) 訓練過程不穩定。

三、核心問題 (The Core Problem)#

基於上述背景，這篇論文要解決的核心問題可以歸結為：

如何在維持或甚至降低單一樣本計算成本 (FLOPS per token) 的前提下，繼續大幅度增加模型的總參數數量，從而突破密集模型的效率瓶頸，並將 MoE 從一個難以駕馭的理論變為一個簡單、穩定且高效的實用架構？

作者們假設，參數數量本身就是一個獨立於總計算量的重要擴展維度。他們要做的，就是找到一把鑰匙，解鎖這個維度的潛力。

四、提出的方法與核心設計#

作者們提出了 Switch Transformer，其設計哲學是「化繁為簡」。它在標準的 Transformer 架構上做了一個看似微小卻至關重要的改動：將 Transformer 中間的 前饋神經網路層 (Feed-Forward Network, FFN) 替換為一個稀疏的 Switch FFN 層。

4.1 傳統 MoE 的數學原理#

要理解 Switch 的創新，我們先要看傳統 MoE 的做法。對於一個輸入 token 的表示 x，和 N 個專家 $\{E_i\}_{i=1}^N$ ，傳統 MoE 的路由過程分兩步：

計算路由概率： 一個可訓練的路由器權重矩陣 $W_r$ 會為所有專家打分，並通過 Softmax 函數轉換為概率分佈 $p(x)$ 。
$p_i(x) = \frac{e^{h(x)_i}}{\sum_{j=1}^{N} e^{h(x)_j}} \quad \text{其中 } h(x) = W_r \cdot x$
這個公式（原論文 Equation 1）計算了 token x 應該被發往第 i 個專家的概率 $p_i(x)$ 。
Top-k 路由與加權輸出： 傳統 MoE 會選擇概率最高的 k 個專家 ( k>1 )，然後將 token x 分別送入這些專家進行計算，最後將它們的輸出 $E_i(x)$ 根據概率進行加權求和，得到最終輸出 y。
$y = \sum_{i \in \mathcal{T}} p_i(x)E_i(x) \quad (\mathcal{T} \text{是 top-k 專家的索引集合})$
這個公式（原論文 Equation 2）描述了最終輸出是如何由多個專家的結果混合而成的。過去的研究者認為 k>1 是必須的，因為這樣才能讓路由器在多個選項之間進行比較和學習。

4.2 Switch Transformer 的簡化與創新#

Switch Transformer 的核心突破，就是推翻了 k>1 的假設，將路由策略極簡化為 k=1。

這意味著，對於每個 token，只選擇概率最高的那個專家。於是，上面的輸出公式就簡化為：

y = p_{top\_1}(x) E_{top\_1}(x)

即輸出就是被選中的那個專家的計算結果，再乘以它的門控值（概率）。

圖表解讀 (Figure 2): 這張圖完美詮釋了 Switch Transformer 的工作流程。左側是標準 Transformer 模塊的示意，右側藍色框內則是創新的 Switch FFN 層。對於輸入的兩個 tokens “More” 和 “Parameters”，路由器（Router）為它們獨立地進行決策，分別選擇了最匹配的專家 FFN 2 和 FFN 3 進行處理（實線箭頭）。其他專家（FFN 1, FFN 4）則完全不參與計算，處於「關閉」狀態。最終的輸出是專家結果與路由器門控值的乘積（虛線箭頭）。

這個 k=1 的設計，帶來了三大好處：路由器計算量減少、專家容量需求降低、通訊成本降低與實現簡化。

五、關鍵概念解析#

5.1 概念概述#

稀疏激活 (Sparsely-Activated): 這是與「密集 (Dense)」相對的概念。密集模型中，所有參數都參與運算。稀疏激活模型中，只有一小部分參數被激活和使用。Switch Transformer 就是典型的稀疏激活模型。
專家 (Expert): 在這篇論文的上下文中，一個專家就是一個標準的前饋神經網路 (FFN) 。模型中有很多個這樣的專家，它們的權重各不相同，因此各有專長。
路由器 (Router): 一個小型的神經網路，功能像個分發員。它判斷每個 token 應該交給哪個專家處理最合適。
專家容量 (Expert Capacity): 由於硬體限制，每個專家能同時處理的 token 數量是有限的，這個上限就是「容量」。如果路由器分配給某個專家的 token 數量超出了其容量，就會發生「溢出」，多餘的 token 會被「丟棄」（通常是直接通過殘差連接跳到下一層，不做任何處理）。
負載平衡損失 (Load Balancing Loss): 為了防止路由器總把任務交給少數幾個「明星專家」，導致忙的忙死、閒的閒死，作者引入了一項輔助損失函數。這個損失函數會懲罰不均衡的分配策略，鼓勵路由器將 token 盡可能均勻地分配給所有專家，從而提高整體訓練效率。

5.2 專家容量 (Expert Capacity)#

由於硬體（如 TPU）要求靜態的張量形狀，每個專家能同時處理的 token 數量必須預先固定，這就是「專家容量」。其計算公式為：

\text{專家容量} = \left( \frac{\text{總 token 數}}{\text{專家數}} \right) \times \text{容量因子 (capacity factor)}

這個公式（原論文 Equation 3）中，「容量因子」是一個大於等於 1.0 的超參數，用作緩衝區。

圖表解讀 (Figure 3): 這張圖生動地展示了容量因子的作用。左側容量因子為 1.0，token 分配不均時，Expert 1 接收了超過其容量的 token，導致「溢出」（紅色虛線），溢出的 token 將不被處理。右側容量因子為 1.5，專家容量更大，可以容納不均勻的分配，但代價是可能產生閒置的計算槽位（白色空格），造成一定的計算浪費。因此，選擇合適的容量因子是在避免 token 丟失和提高計算效率之間的一種權衡。

5.3 負載平衡損失 (Load Balancing Loss)#

為了鼓勵路由器將 token 均勻分配，作者引入了一項輔助損失函數。對於一個批次的數據，該損失的計算方式為：

\text{loss} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i

這個公式（原論文 Equation 4）是整個負載平衡機制的核心。其中：

$N$ 是專家總數。
$\alpha$ 是一個超參數，用於調節該損失項的重要性。
$f_i$ 是實際被路由到專家 i 的 token 數量佔總數的比例。
$P_i$ 是路由器預測的概率分配給專家 i 的平均比例。

理想情況下，token 應該被均勻分配，此時對於任何專家 i， $f_i$ 和 $P_i$ 都應該趨近於 $1/N$ 。這個公式通過最小化兩個向量的點積，來促使路由器的分配策略趨向於均勻分佈。

六、實驗設計與結果分析 (圖表證據)#

論文通過大量實驗，從多個角度展示了 Switch Transformer 的壓倒性優勢。

無可匹敵的訓練速度：

圖表解讀 (Figure 5): 這是最有說服力的圖表之一。它比較了在相同計算資源下，不同模型達到相同性能（負對數困惑度 Neg Log Perplexity）所需的實際訓練時間。可以看到，Switch-Base 模型（如綠線和橙線）的學習曲線比密集模型 T5-Base（紅線）陡峭得多。Switch-Base (64個專家) 達到 T5-Base 最終性能水平，只用了大約 1/7 的時間。這證明了其在時間效率上的巨大優勢。
卓越的擴展性與樣本效率：
圖表解讀 (Figure 4): 這張圖揭示了 Switch Transformer 的擴展潛力。
- 左圖顯示，在計算量 (FLOPS) 固定的情況下，隨著專家數量從 1 增加到 256，模型的總參數從數億增加到近 150 億，而測試損失 (Test Loss) 持續穩定下降。這證明了增加稀疏參數的有效性。
- 右圖比較了不同專家數量的 Switch-Base 模型與 T5-Base 的樣本效率。所有 Switch 模型的學習速度（每一步的性能提升）都遠超 T5-Base。
碾壓更大規模的密集模型：

圖表解讀 (Figure 6): 為了回答「與其用 Switch，不如直接訓練一個更大的密集模型？」這個問題，作者將 Switch-Base 與計算量是其 3.5 倍的 T5-Large 進行了比較。結果顯示（右圖），即便是面對這個更強大的密集對手，Switch-Base 依然在訓練速度上取得了 2.5 倍的領先優勢。這說明，稀疏化是比單純增加密集計算量更高效的擴展路徑。
關鍵的訓練穩定性技術： 論文通過實驗驗證了多項穩定性技術的有效性。例如，Table 2 的數據顯示，「選擇性精度」訓練在與不穩定的 bfloat16 保持幾乎相同速度的同時，達到了與穩定的 float32 相媲美的性能。

Table 3 則證明了減小初始化尺度能顯著提升模型早期的平均性能並降低不穩定性。
萬億參數模型的實現： 論文成功訓練了 高達 1.6 萬億參數 的 Switch-C 模型，並在與 Google 自家的 110 億參數的密集模型 T5-XXL 的對比中，取得了 4 倍的訓練速度提升，首次向世界展示了萬億級參數模型在訓練效率上的可行性。
廣泛的有效性： 無論是在下游任務微調、多語言環境，還是將大模型蒸餾壓縮成小模型，Switch Transformer 都表現出色。

七、論文的核心貢獻與意義#

重新定義了效率： 它證明了模型的強大不僅僅來源於計算量 (FLOPS)，更來源於其巨大的「知識容量」（參數數量）。通過「化繁為簡」的 Switch 路由機制，它開創了一條 計算成本可控、參數規模無限 的新 scaling law 路徑。
讓稀疏模型成為主流： 它掃除了一直以來阻礙 MoE 發展的複雜性、通訊和穩定性障礙，提供了一套完整、可行的解決方案（如 k=1 路由、負載平衡、選擇性精度等），使得稀疏架構成為了大型模型設計的一線選擇。
突破了萬億參數的壁壘： 論文成功訓練了高達 1.6 萬億參數的 Switch-C 模型，並在與 T5-XXL 的對比中取得 4 倍速度提升，首次在實踐中證明了訓練如此規模的模型不僅是可能的，而且可以非常高效。