專家混合（MoE）說明

專家混合（MoE）使 transformer 模型的預訓練速度遠快於密集模型，且相較於總參數量，推論速度也更快。這是透過將密集前饋網路（FFN）層替換為稀疏 MoE 層來實現的，稀疏 MoE 層在任意輸入 token 上僅啟用少數「專家」。

專家混合的核心架構

MoE 模型以稀疏 MoE 層與路由機制的組合取代標準 FFN 層。此架構包含兩個主要組件：

  • 稀疏 MoE 層：這些層不是單一的密集層，而是包含多個「專家」，每個專家通常是一個神經網路（常為 FFN）。這些專家可以以階層結構排列，專家本身也可以是 MoE。
  • 門控網路（路由器）：一個學習得到的路由器決定哪些 token 會被送至哪個專家。路由器與其餘網路一起預訓練，且可設定為將 token 送至一個或多個專家。

雖然 MoE 提供計算效率，但也帶來特定的取捨。它們需要大量 VRAM，因為所有專家必須載入記憶體，即使在單次前向傳播時僅有少數專家被啟用。例如，Mixtral 8x7B 需要足夠的 VRAM 以容納 47B 參數的模型，儘管在每個 token 使用兩個專家時，其推論計算量（FLOPs）接近 12B 模型。

稀疏性與路由機制

稀疏性允許模型在規模擴大時，透過條件計算——僅在每個樣本上啟用系統的部分——而不會成比例增加計算量。

閘控函數

路由由閘控網路 $G$ 處理，決定要啟用哪些專家 $E$。傳統設定使用 softmax 函數。更進階的方法，例如 Noisy Top-k Gating，會加入可調噪聲，並選取前 $k$ 個值，以確保訓練期間有更好的探索與負載平衡。

負載平衡與專家容量

為防止「贏者通吃」的情況，即少數專家被過度使用而其他專家被忽視，MoE 採用：

  • 輔助損失：額外的損失函數，鼓勵路由器將 token 均勻分配至所有專家。
  • 專家容量：限制單一專家可處理的 token 數量的閾值。若專家達到容量，token 可能被丟棄或透過殘差連接傳遞至下一層。

MoE 實作的演進

GShard 與 Switch Transformers

Google 的 GShard 透過使用 top-2 閘控與引入隨機路由（第二個專家依其權重比例被選取），將 transformer 的規模擴展至超過 6000 億參數。

Switch Transformers 進一步簡化此概念，採用單一專家策略。每個 token 只路由至一個專家，使 Switch Transformers 減少路由計算、將專家批次大小減半，並降低通信成本，同時保持品質。此架構相較於 T5-XXL 提供了 4 倍的預訓練加速。

穩定性與專家專精

訓練穩定性透過如 Router z-loss（在 ST-MoE 中引入）等技術來管理，該技術對進入閘控網路的大 logits 施加懲罰，以減少捨入誤差。

對專家專精的研究顯示，編碼器的專家常在淺層概念（例如標點或專有名詞）上專精，而解碼器的專家則較少專精。在多語言設定中，因負載平衡需求，專家通常不會僅專精於單一語言。

微調與指令微調

稀疏模型較密集模型更容易過擬合。有效的微調策略包括：

  • 較高正則化：在稀疏層中使用較高的 dropout 率。
  • 選擇性凍結：凍結 MoE 層，同時更新非 MoE 層，可在加速訓練的同時保持品質。
  • 超參數調整：稀疏模型通常受惠於較小的批次大小與較高的學習率。

近期研究顯示，MoE 從 指令微調 中獲得的效益遠高於密集模型。將 MoE 與等效的 T5 進行比較時，MoE 透過指令微調（Flan-MoE vs MoE）所得到的性能提升，超過了密集模型的提升（Flan T5 vs T5）。

計算效率與平行化

平行化策略

為處理 MoE 巨大的參數量，使用了不同的平行化技術：

  • 專家平行化：將專家分散至不同的工作節點。對於 MoE 層，token 會被路由至該專家所在的工作節點；對於非 MoE 層，則類似資料平行化。

硬體最佳化

  • 容量因子：在品質與通信成本之間的取捨。較高的容量因子提升品質，但會增加激活的記憶體需求與設備間的通信。
  • 區塊稀疏運算：如 MegaBlocks 等專案以區塊稀疏運算取代傳統的批次矩陣乘法，以處理不平衡的 token 分配而不丟棄 token，從而顯著加速。

比較：稀疏 MoE 與 密集模型

特徵 稀疏 MoE 密集模型
預訓練計算 更有效率 / 更快 較低效率
推論速度 更快（相對於總參數） 較慢（相對於總參數）
VRAM 需求 高（必須載入所有專家） 較低（相對於總參數）
最佳使用情境 高吞吐量，多機器 低吞吐量，VRAM 受限

開源 MoE 生態系統

多個框架與模型將 MoE 帶入開源社群：

  • 框架：MegaBlocks、Fairseq 與 OpenMoE。
  • 模型：Google 的 Switch Transformers（最高 1.6T 參數）、Meta 的 NLLB MoE（54B）以及 Mistral 的 Mixtral 8x7B，該模型在推論速度上優於 Llama 2 70B。

Sources