在標準的 Transformer 中,每個 token 都通過每一層中相同的前饋網路(FFN)。在 MoE Transformer 中,這個單一的 FFN 被替換為多個平行的 FFN —— 即「專家」—— 再加上一個小型路由網路(通常稱為閘門),由它決定哪些專家處理每個 token。通常,閘門會選出 top-k 個專家(一般是 2 個),並使用閘門的 softmax 權重混合它們的輸出。其關鍵洞見在於:總參數量可以非常龐大(讓模型擁有巨大的記憶和泛化能力),而每個 token 的計算量卻可控,因為大多數專家對任何給定的輸入都是閒置的。例如,Mixtral 8x7B 總共約有 470 億參數,但每個 token 只啟動約 130 億。
路由機制是工程複雜度最集中的地方。一個簡陋的路由器可能會把所有 token 都送到同幾個專家那裡,讓其他專家閒置 —— 這個問題稱為「專家崩潰」。為了避免這種情況,MoE 模型使用輔助的負載平衡損失函數,在訓練過程中懲罰不均勻的專家使用率。Google 原始的 Switch Transformer 使用 top-1 路由(每個 token 一個專家),並實現了出色的擴展性,但大多數現代 MoE 模型偏好 top-2 路由以提高穩定性。一些較新的方法如 DeepSeekMoE 則加入了「共享專家」,這些專家始終與路由選出的專家一起啟動,確保每個 token 無論路由決策如何都能獲得基本水準的處理。
定義 MoE 部署的核心權衡是記憶體 vs. 計算。即使每個 token 只有少數專家處於活躍狀態,所有專家都必須載入記憶體。一個 8x7B 的 MoE 模型所需的記憶體大約等同於一個稠密的 47B 模型,即使它的運行速度大約等同於一個稠密的 13B 模型。這使得 MoE 模型在消費級硬體上有些尷尬 —— 如果你的 GPU VRAM 只能容納 13B 參數,用一個稠密的 13B 模型可以獲得同樣的推理速度,卻沒有 MoE 的額外開銷。MoE 真正發揮優勢的場景是當你有足夠的記憶體容納完整模型,並且希望在每次浮點運算中獲得最大品質時。這就是為什麼它天生適合雲端服務:像 OpenAI 和 Mistral 這樣的供應商可以在叢集上配置足夠的記憶體,而每次請求的計算成本才是影響他們利潤的關鍵。
專家並行是一種 MoE 特有的部署模式。在多 GPU 設定中,你可以將不同的專家放置在不同的 GPU 上,這樣每台設備只需儲存和計算專家的一個子集。Token 根據所需的專家被路由到不同的 GPU,處理後再匯集結果。這會引入 all-to-all 的通訊開銷,但允許遠超單一設備容量的模型高效運行。Google 的 GShard 和 Switch Transformer 論文大規模展示了這種方法,這也是當今最大的 MoE 模型在生產環境中的運作方式。
實務工作者會遇到的一個問題是:MoE 模型在微調時可能表現不穩定。如果你在一個狹窄的領域上進行微調,路由器可能會開始將所有 token 導向少數幾個專家,實質上浪費了其餘專家的能力。一些團隊在微調時凍結路由器;另一些則加入額外的正則化。量化是另一個陷阱 —— 很少被啟動的專家校準樣本較少,因此簡單的訓練後量化會對它們造成不成比例的性能下降。業界正在積極解決這些實務挑戰,但 MoE 顯然是未來的方向。Grok、DBRX、Mixtral,以及幾乎可以確定的 GPT-4 都使用了它,而隨著模型規模增長,效率上的論點只會越來越有力。