Switch Transformer

Switch Transformer 深度學習頁

3 秒看懂

Switch Transformer = 每個 token 只路由到 1 個專家（top-1 gating）的稀疏 MoE 架構，用更少的計算量撬動萬億級引數量，實現訓練速度數量級躍升。

一句話定位：它不是”更大的 Transformer”，而是”用稀疏啟用讓引數量與計算量脫鉤”的範式開創者。

3 分鐘產業解釋

它解決什麼問題？

傳統 Dense Transformer 的引數量和計算量（FLOPs）是線性繫結的——模型大一倍，推論/訓練計算也大一倍。這使得超大型模型的訓練成本呈平方級攀升。

Switch Transformer 的核心洞察：如果我們允許每個 token 只啟用模型中一小部分引數（稀疏啟用），就可以將引數量推到萬億級別，而每個 token 的實際計算量仍然可控。

類比

想像一所擁有 128 個科室的大型醫院（128 個專家）。傳統”稠密”模型要求每位患者同時看所有 128 個科室——荒謬且昂貴。Switch Transformer 的做法是：每位患者進入大廳後，前臺（路由器）根據症狀直接分配到唯一一個最匹配的科室。單個患者花的錢少了，但醫院總體接待能力大幅增加。

產業意義

Google Brain 團隊 2021 年釋出，論文發表於 JMLR 2022
首次系統驗證了 top-1 路由在 MoE 中的可行性與效率優勢
直接影響了後續一系列稀疏/混合架構：GShard → Switch → ST-MoE → Mixtral → DeepSeek-MoE
論文標題本身即宣言：“Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity”

15 分鐘專家深入

核心架構改動

Switch Transformer 的改動極其剋制——僅將 Transformer 的 FFN（前饋網路）層替換為 MoE 層，其餘結構（注意力、歸一化、殘差連線）完全不變。

標準 Transformer Block:
  Input → LayerNorm → Self-Attention → Residual → LayerNorm → FFN → Residual → Output

Switch Transformer Block:
  Input → LayerNorm → Self-Attention → Residual → LayerNorm → MoE-FFN → Residual → Output
                                                      ↑
                                              路由器 + N 個專家FFN

Top-1 路由機制

這是 Switch Transformer 最核心的貢獻。設輸入 token 的隱藏狀態為 \mathbf{x}，路由器是一個簡單的線性層加 softmax：

G(\mathbf{x}) = \text{Softmax}(\mathbf{W}_r \cdot \mathbf{x})

其中 \mathbf{W}_r \in \mathbb{R}^{N \times d}，$N$ 為專家數，$d$ 為隱藏維度。

關鍵決策：取 \text{argmax}，選擇機率最高的唯一一個專家：

i^* = \arg\max_i G(\mathbf{x})_i

\text{MoE-FFN}(\mathbf{x}) = G(\mathbf{x})_{i^*} \cdot \text{FFN}_{i^*}(\mathbf{x})

注意：輸出仍乘以門控權重 G(\mathbf{x})_{i^*}，這是一個連續值，保留了專家輸出的”置信度”縮放。

與 Top-2 路由的關鍵區別

維度	Top-2（GShard 等）	Top-1（Switch）
每 token 啟用專家數	2	1
FFN 前向計算量/token	2× 單專家 FFN	1× 單專家 FFN
All-to-All 通訊量	較高（需分發到 2 個專家）	較低
實現複雜度	需要 combine 步驟	更簡單
質量	略優	接近（損失很小）

論文實驗證實：top-1 在相同計算預算下效能接近 top-2，但吞吐量顯著更高——這是工程上的關鍵勝利。

容量因子（Capacity Factor）

這是防止負載不均衡導致 token 被丟棄的關鍵超引數。

每個專家被分配一個 buffer 容量 $C$，定義為： C = \text{CF} \times \frac{T}{N} 其中 $T$ 是序列中的 token 總數，$N$ 是專家數，\text{CF} 是容量因子。
CF = 1.0：每個專家恰好分到平均數量的 token（理想均勻時剛好夠）
CF > 1.0（如 1.2、1.5）：給每個專家額外餘量，容納不均勻分配
CF < 1.0：強制丟棄部分 token（token dropping），節省計算但損失資訊

工程權衡：CF 越大越穩定但計算浪費越多；CF 越小越高效但丟 token 的風險越高。論文中推薦 CF 在 1.0–1.5 範圍內調節。

輔助負載均衡損失

為避免”贏者通吃”（少數專家被過度使用，其餘閒置），Switch Transformer 引入輔助損失 \mathcal{L}_{\text{aux}}：

\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i

其中：

f_i = 被路由到專家 $i$ 的 token 比例
P_i = 路由器分配給專家 $i$ 的平均機率
\alpha = 輔助損失係數（論文中常設為 10^{-2} 量級）
$N$ = 專家數量

該項在 f_i = P_i = 1/N（完美均勻）時取最小值，激勵均勻分配。

選擇性精度訓練（Selective Precision）

一個關鍵的工程細節：路由器中的 softmax 操作必須在 fp32 精度下進行，即使模型其餘部分使用 bf16/fp16 混合精度。路由器的線性變換 \mathbf{W}_r \cdot \mathbf{x} 通常可以保留在 bf16 計算，僅 softmax 提升至 fp32 以保證數值穩定性。

原因：softmax 輸出是細粒度機率分佈，低精度下容易出現數值不穩定（如 softmax 飽和、梯度消失），導致路由坍縮到少數專家。

┌─────────────────────────────────────┐
│         Switch Transformer Block     │
│                                      │
│  Token x (bf16)                      │
│       │                              │
│       ├──→ Expert FFNs (bf16)        │
│       │                              │
│       └──→ Router W_r (bf16)         │
│              │                       │
│              ↓                       │
│         Softmax (fp32)               │
│              │                       │
│              ↓                       │
│         argmax → 選擇專家 i*         │
│              │                       │
│              ↓                       │
│         輸出 = gate × FFN_i*(x)      │
└─────────────────────────────────────┘

技術原理（最深一層）

完整前向傳播流程

輸入序列: x ∈ R^&#123;seq_len × d_model&#125;
         │
         ▼
    ┌─ Self-Attention ─┐
    │  (標準 MHA/MQA)   │
    └────────┬─────────┘
             │ residual
             ▼
         LayerNorm
             │
             ▼
    ┌─ Switch MoE Layer ─────────────────────────┐
    │                                              │
    │  1. 路由計算:                                │
    │     S = Softmax(x · W_r^T)  ∈ R^&#123;seq_len×N&#125;│
    │     (W_r ∈ R^&#123;N × d&#125;, fp32)                 │
    │                                              │
    │  2. Top-1 分派:                              │
    │     expert_id = argmax(S, dim=-1)            │
    │     gate = max(S, dim=-1)                    │
    │                                              │
    │  3. Token 重排 (Permutation):                │
    │     按 expert_id 對 token 分組/排序           │
    │     形成 N 個 token bucket                    │
    │                                              │
    │  4. 容量截斷:                                │
    │     每個 bucket 只保留前 C 個 token           │
    │     超出部分被丟棄 (dropped tokens)           │
    │                                              │
    │  5. 並行專家計算:                            │
    │     每個專家 FFN_i 獨立處理其 bucket          │
    │     FFN_i: x → W_up · GELU(x) · W_down      │
    │     (bf16, 與標準 Transformer FFN 結構相同)  │
    │                                              │
    │  6. 門控縮放 + 反排列:                       │
    │     output_j = gate_j × FFN_&#123;id_j}(x_j)     │
    │     恢復原始 token 順序                       │
    │                                              │
    └──────────────────────────────────────────────┘
             │ residual
             ▼
          輸出

通訊模式

在分散式訓練中，MoE 層引入 All-to-All 通訊：

Step 1 - Dispatch (All-to-All):
  Device 0 持有的 token → 需要傳送給 Device 1, 2, 3 上的專家
  Device 1 持有的 token → 需要傳送給 Device 0, 2, 3 上的專家
  ...以此類推

Step 2 - Expert Computation:
  每個裝置上的專家處理收到的 token

Step 3 - Combine (All-to-All):
  處理結果傳送回原始裝置

通訊量與專家並行度和 token 數量成正比。Switch Transformer 的 top-1 路由相比 top-2 路由，每個 token 只需訪問一個裝置的專家，因此 combine 通訊量減半。

關鍵引數量級估算

以論文中描述的 Switch-Base 配置為參考（注意：以下為基於論文描述的定性估算，具體數值請參閱原論文 Table）：

隱藏維度 $d$：768（與 T5-Base 對齊）
注意力頭數：12
層數：12（encoder）+ 12（decoder）
專家數 $N$：典型為 128（Switch-Base-128 變體）
單個專家 FFN 引數量 ≈ 標準 T5-Base FFN 引數量 / 128 × 128… 等等，這裡需要澄清

澄清：專家數量增加時，每個專家的 FFN 大小不變（與原始 Dense FFN 相同），因此 MoE 層的 FFN 總引數量 = N × 單個 FFN 引數量。模型總引數量因專家數線性增長，但每個 token 的 FLOPs 只取決於 1 個 FFN（top-1）。

論文中最大的變體引數量達到萬億級別（~1.6T），專家數為 2048，但每個 token 的計算量僅相當於一個較小的 Dense 模型。

技術演進史

時間	里程碑	關鍵進展
2017	Shazeer et al. Outrageously Large Neural Networks	首次將 MoE 引入 LSTM/GRU，top-k 路由 + 噪聲探索
2020	GShard（Lepikhin et al.）	MoE 引入 Transformer，encoder 中使用 top-2 路由，擴充套件到 600B 引數
2021	Switch Transformer（Fedus, Zoph, Shazeer）	Top-1 路由、簡化 MoE、系統性 scaling 研究
2021	MoE 分散式訓練最佳化（Various）	Megablocks（Stanford/Databricks）最佳化稀疏計算 kernel
2022	ST-MoE（Zoph et al., Google）	穩定訓練的 MoE（encoder-decoder），訓練穩定性與遷移能力研究
2023	Mixtral 8×7B（Mistral AI）	MoE 落地開源大語言模型，decoder-only，實際部署
2024	DeepSeek-MoE / DeepSeek-V2	共享專家 + 細粒度路由專家，進一步最佳化專家利用率
2024	Mixtral 8×22B、DBRX（Databricks）	MoE 架構成為開源大型模型的主流路線之一

技術路線對比

維度	Dense Transformer (T5/GPT)	Switch Transformer	DeepSeek-V2 MoE	Mixtral 8×7B
路由策略	N/A（全引數啟用）	Top-1, 所有 token	Top-K + 共享專家	Top-2
每 token FLOPs	全部 FFN 引數	1/N × FFN 引數	可配置	2/N × FFN 引數
總引數量級	受計算預算約束	可大幅超越計算預算	同左	同左
通訊模式	AllReduce（張量並行）	All-to-All（專家並行）	All-to-All	All-to-All
訓練穩定性	成熟	需要選擇性精度 + 輔助損失	進一步改進	相對穩定
推論友好度	高（引數全部駐留視訊記憶體）	低（需 All-to-All）	中等	中等
適用架構	Encoder / Decoder / Enc-Dec	Enc-Dec（原論文）	Decoder-only	Decoder-only

上下游

上游依賴

資料與基礎設施層
├── 訓練資料：C4（Colossal Clean Crawled Corpus）等大規模語料
├── 訓算硬體：Google TPU（論文基於 TPU Pod 訓練）
├── 架構：Mesh TensorFlow / JAX（Google 內部生態）
└── 基礎架構：T5 Encoder-Decoder 骨架

關鍵技術輸入層
├── MoE 路由理論：Shazeer 2017, GShard 2020
├── Transformer 架構：Vaswani et al. 2017
├── 稀疏計算思想：條件計算 (Conditional Computation)
└── 負載均衡理論：最優傳輸 / 均勻分配

下游影響

直接後繼架構
├── ST-MoE (Google, 2022) —— 穩定化訓練
├── Mixtral 8×7B / 8×22B (Mistral, 2023-2024) —— 開源 MoE LLM
├── DeepSeek-MoE / V2 / V3 (DeepSeek, 2024-2025) —— 細粒度路由
├── DBRX (Databricks, 2024) —— 企業級 MoE
└── Qwen-MoE / Grok-1 等

工程生態影響
├── Megablocks (Databricks) —— 稀疏計算 CUDA kernel
├── Tutel (Microsoft) —— MoE 通訊最佳化
├── vLLM / TensorRT-LLM —— MoE 推論最佳化
└── HuggingFace Transformers —— Switch Transformer 整合

關鍵指標

指標	含義	Switch Transformer 典型值/範圍
專家數 $N$	MoE 層中並行 FFN 的數量	8–2048（論文覆蓋範圍）
容量因子 CF	每個專家 buffer 相對平均分配的倍數	1.0–1.5
輔助損失係數 `\alpha`	負載均衡損失的權重	~0.01 量級
Token dropping rate	因容量不足被丟棄的 token 比例	取決於 CF 和均勻度
FLOPs 相對加速	相同質量下相比 Dense 的訓練速度	論文報告約 7×（Switch-Base vs T5-Base）
引數效率比	總引數量 / 每 token 啟用引數量	N 倍（N=專家數）

供需與市場資料

⚠️ 注意：Switch Transformer 本身是 Google Research 的學術成果，無獨立商業化產品。以下為 MoE 架構在產業中的市場對映。

MoE 模型的訓練計算需求（估算）

模型	架構	總引數量	訓練硬體（公開資訊）
Switch-C（最大變體）	Enc-Dec MoE	~1.6T [論文]	Google TPU Pod（具體規格未充分揭露）
Mixtral 8×7B	Dec-only MoE	~46.7B（啟用 ~12.9B）[Mistral]	未充分揭露，估算數百至千卡叢集
DeepSeek-V3	Dec-only MoE	~671B（啟用 ~37B）[DeepSeek]	2048 H800 [DeepSeek 技術報告]
Grok-1	Dec-only MoE	~314B（公開）[xAI]	未充分揭露

市場趨勢判斷

MoE 架構已從學術概念快速過渡為產業主流：

訓練側：MoE 降低同等質量下的訓練 FLOPs，但引入通訊開銷，對網路頻寬要求更高
推論側：MoE 模型的總引數量大但啟用引數少，視訊記憶體佔用仍需容納全部專家引數，但每 token 計算量低——形成”視訊記憶體牆 + 計算效率”的獨特約束
硬體適配：MoE 推動了對高頻寬互聯（NVLink/NVSwitch, TPU ICI）和大視訊記憶體（HBM3/3e）的需求

代表公司與資本對映

公司/機構	與 Switch Transformer / MoE 的關係	資本對映
Google (Alphabet)	Switch Transformer 原研團隊所在；PaLM-2 使用 MoE 變體	GOOGL
Mistral AI	Mixtral 系列——最成功的開源 MoE LLM	未上市（估值~$6B, 2024）
DeepSeek (幻方量化)	DeepSeek-V2/V3——創新的 MoE 路由設計	非上市（母公司幻方）
NVIDIA	MoE 模型訓練和推論高度依賴 GPU + NVLink 互聯	NVDA
Databricks	DBRX MoE 模型 + Megablocks 開源架構	未上市（估值~$43B, 2023）
xAI	Grok-1 採用 MoE 架構	未上市
HuggingFace	Transformers 庫整合 Switch/Mixtral 等 MoE 模型	未上市

投資邏輯

核心論點

MoE 已成為大型模型 Scaling 的確定性路線之一
- Switch Transformer 證明了稀疏啟用的可行性，Mixtral/DeepSeek-V3 驗證了產業落地性
- 在同等訓練預算下，MoE 能達到更高質量——這意味著算力效率提升 = 成本下降 = 誰先掌握誰受益
MoE 對基礎設施的需求是”不同”而非”更少”
- 訓練 MoE 需要更高頻寬的互聯（All-to-All 通訊模式）→ 利好 NVLink/NVSwitch 生態
- 推論 MoE 需要更大視訊記憶體容納全部專家引數 → 利好 HBM 供應商
- MoE 不是”省算力”，而是”重塑算力結構”
開源 MoE 模型降低准入門檻
- Mixtral 8×7B 在普通 A100 叢集上可部署 → 加速下游應用爆發
- DeepSeek 開源權重 → 中國 AI 生態的 MoE 研究加速

風險點

MoE 推論的通訊開銷使其在邊緣/端側部署困難（當前主流仍是 Dense 小模型）
負載不均衡問題在實際生產中仍是未完全解決的工程挑戰
如果硬體架構（如存算一體）根本性改變計算範式，MoE 的優勢可能被重新評估

常見誤讀糾偏

❌ 誤讀 1：“Switch Transformer 的每個專家引數更少，所以模型更小”

糾正：這是最常見的誤解。Switch Transformer 的總引數量 = N × 單個專家FFN引數量 + 共享引數（注意力層等），遠大於同尺寸 Dense 模型。每個 token 只啟用一個專家，所以單 token 計算量小，但模型總引數量可以非常大（論文中到 1.6T）。區分清楚：啟用引數 ≠ 總引數。

❌ 誤讀 2：“Switch Transformer 的 top-1 路由是 GShard top-2 的簡化版，效果更差”

糾正：論文實驗表明，在相同計算預算（FLOPs）下，top-1 路由的訓練效率接近甚至有時優於 top-2。top-2 確實在某些指標上有微弱優勢，但 top-1 的吞吐量顯著更高（只需一半的 FFN 計算和更少的通訊），綜合收益更大。這不是”簡化降級”，而是”減法設計的工程勝利”。

❌ 誤讀 3：“MoE 模型推論更快因為啟用引數少”

糾正：部分正確但容易誤導。MoE 模型的每 token FLOPs 確實更少（只啟用 1/N 個專家的 FFN），但：

推論時仍需將全部專家引數載入到視訊記憶體中
All-to-All 通訊引入延遲（尤其在多 GPU 推論時）
批處理效率可能因 token 路由到不同專家而降低
實際吞吐量提升取決於具體實現、批大小、硬體拓撲

❌ 誤讀 4：“Switch Transformer 是 encoder-decoder，跟現在的 LLM 沒關係”

糾正：Switch Transformer 的核心機制（top-1 路由、容量因子、輔助損失、選擇性精度）是架構無關的。後續 Mixtral、DeepSeek-V2/V3、Grok-1 等 decoder-only LLM 全面採用了 Switch 提出的稀疏 MoE 範式。encoder-decoder 只是原論文選擇的實驗平台。

學習路徑

入門（1–2 小時）

閱讀論文 Abstract + Introduction + Section 3（Switch Routing）—— arXiv:2101.03961
HuggingFace 部落格：“Mixture of Experts Explained”（視覺化講解 MoE 基本概念）

進階（半天）

論文完整閱讀，重點關注 Section 4（訓練穩定性）和 Section 5（Scaling Results）
對照閱讀 GShard 論文（Lepikhin et al., 2020），理解 top-2 vs top-1 的設計取捨
閱讀 ST-MoE 論文（Zoph et al., 2022），瞭解後續改進

工程實踐

HuggingFace Transformers 中執行 Switch Transformer 示例程式碼（switch-base-8 等小變體）
閱讀 Megablocks 程式碼庫（Databricks），理解稀疏計算的 CUDA kernel 實現
閱讀 Mixtral 8×7B 技術部落格，理解 MoE 在 decoder-only LLM 中的實際部署

前沿追蹤

DeepSeek-V2 技術報告——理解共享專家 + 細粒度路由的創新
DeepSeek-V3 技術報告——瞭解大規模 MoE 的工程實踐（負載均衡策略、通訊最佳化等）
關注 MoE 推論最佳化：vLLM 的 Expert Parallelism 支援、TensorRT-LLM 的 MoE 最佳化

一句話總結

Switch Transformer 用”每個 token 只走一條路”的極簡路由，證明了稀疏 MoE 可以高效擴充套件到萬億引數，奠定了此後所有 MoE 大型模型的工程範式基礎。

延伸閱讀與來源

一手論文

Fedus, Zoph, Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR, 2022. arXiv:2101.03961
Shazeer et al. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017. arXiv:1701.06538
Lepikhin et al. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ICLR 2021. arXiv:2006.16668

後續重要工作

Zoph et al. ST-MoE: Designing Stable and Transferable Sparse Expert Models. 2022. arXiv:2202.08906
Jiang et al. Mixtral of Experts. 2024. arXiv:2401.04088
DeepSeek-AI. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. 2024. arXiv:2405.04434

工程資源

HuggingFace Transformers MoE 文件：https://huggingface.co/docs/transformers/moe
Megablocks (Databricks)：https://github.com/databricks/megablocks
Tutel (Microsoft)：https://github

Switch Transformer