在混合專家模型（MoE）的實踐中，負載不均衡儼然已成為制約模型性能提升的關(guān)鍵瓶頸之一。傳統(tǒng)的均衡策略往往需要引入復雜的輔助損失函數(shù)，不僅增加了訓練的復雜度，還可能干擾模型的核心學習目標。工程師們在提升模型效率的道路上，一直苦苦追尋著一個優(yōu)雅而高效的平衡解決方案。

DeepSeek 團隊的這項研究，為這一長期困擾業(yè)界的技術(shù)難題提供了令人耳目一新的解決思路：通過在門控分數(shù)中直接添加專家層面的偏置項，在絕大部分不引入額外損失函數(shù)的情況下，實現(xiàn)了模型訓練過程中的自適應負載均衡。更令人驚嘆的是，這一方法不僅保持了模型的因果關(guān)系，還顯著提升了訓練的穩(wěn)定性和最終性能。

我們將探討 DeepSeek 模型[1, 2, 3]中與混合專家模型（MoE）相關(guān)的另一項關(guān)鍵架構(gòu)突破：無輔助損失函數(shù)的負載均衡（Auxiliary-Loss-Free Load Balancing）[5]。

在本文，我們將深入解析 DeepSeek 如何解決 MoE 的隱藏瓶頸——負載均衡，同時還通過消除梯度干擾和嚴格遵循因果關(guān)系約束，提升了訓練和推理效率，為后續(xù)專家模型的優(yōu)化方向提供了標桿。

本文目錄：

• 技術(shù)背景：介紹混合專家模型（MoE）的基本原理，解釋負載均衡的重要性，回顧之前的工作，包括輔助損失函數(shù)（auxiliary loss）方法和專家選擇（Expert Choice）策略。
• DeepSeek 的無輔助損失函數(shù)的負載均衡：解析其運作原理
• 性能評估：討論無輔助損失函數(shù)的負載均衡的性能表現(xiàn)
• 總結(jié)
• 參考文獻

「DeepSeek-V3 技術(shù)解析」系列其他文章：

• 「DeepSeek-V3 技術(shù)解析」：多頭潛在注意力機制（MLA）
• 「DeepSeek-V3 技術(shù)解析」：DeepSeekMoE

1.技術(shù)背景

1.1 MoE (Mixture-of-Experts) in Transformers

MoE（Mixture-of-Experts，混合專家模型）在 Transformer 中的實現(xiàn)方式通常為：每隔若干 Transformer 層，將其中的 FFN（前饋網(wǎng)絡）替換為多個 FFN（每個 FFN 充當一個"專家"）。當 input token 進入該層時，通過門控操作（Gating）選擇 Top-K 個專家，并將 input token 只路由至這些被選中的 FFN，從而僅激活對應的專家網(wǎng)絡。

下圖展示了這一過程：左側(cè)標準 Transformer 層中的 FFN 子層被替換為右側(cè)的 MoE 層。

如需更詳細的 MoE 技術(shù)解析，可參考「DeepSeek-V3 技術(shù)解析」：DeepSeekMoE（文中通過餐廳類比直觀解釋了 MoE 原理）。

圖 1. Transformer 中的 MoE 層（紅框內(nèi)）示意圖。圖片改編自文獻 [6]。

1.2 負載均衡及其重要性

我們通過一個能夠提供多種菜系菜品的餐廳解釋了 MoE 的概念：餐廳的每位廚師都是專家，主廚（Head Chef）的工作類似于門控操作，將每道菜品分配給具備對應技能的特定廚師。

為確保該系統(tǒng)高效運作，需滿足以下條件：

• 每位專業(yè)廚師必須精通自身菜系所需技能（例如餃子廚師必須會包餃子），同時所有廚師需能共同處理所有菜品。
• 主廚需充分了解所有專業(yè)廚師的專長，并能高效分配訂單。

在 MoE 中，前者對應 expert specialization 與 knowledge sharing 的權(quán)衡（我們已在介紹 DeepSeekMoE 的這篇文章中進行了詳細討論），后者則體現(xiàn)了負載均衡的重要性 —— 這也是本文的核心主題。

為何負載均衡如此重要？

原因在于，當負載不均衡發(fā)生時，MoE 無法有效運作。最常見的問題是路由崩潰（Route Collapse）：少數(shù)專家接收大部分 input token，而其他專家的利用率極低。

因此，大部分計算由超負荷工作的專家承擔，而這些專家通常分布在多個 GPU 核心上，因此會導致硬件資源浪費。

由于梯度沖突（gradient conflict），路由奔潰也會導致訓練不穩(wěn)定。超負荷工作的專家接收更多 input token，他們積累的梯度也會更大，學習速度也會比工作負荷不足的專家快得多；因此，兩者的梯度在幅值和方向上均可能發(fā)生偏離，導致訓練難以收斂。

最后，MoE 中的負載不均衡也會導致性能低下和泛化效果不佳，因為工作負荷不足的專家會因為訓練 tokens 不足，而難以學習有效知識。

由于負載均衡技術(shù)對 MoE 至關(guān)重要，因此研究者們針對這一問題提出了多種解決方案。其中，最常用的策略是為負載均衡添加輔助損失函數(shù)（Auxiliary Loss）和專家選擇（Expert Choice）。

1.3 帶輔助損失函數(shù)的負載均衡

一種常見的改善負載均衡的策略是在模型訓練的目標函數(shù)基礎(chǔ)上引入輔助損失函數(shù)。

圖 2. 用于強化負載均衡的輔助損失函數(shù)示例。圖片編輯自文獻[5]。

上圖展示了一個輔助損失函數(shù)的示例，其中

• N 是專家數(shù)量，T 是 token 數(shù)量，K 是每個 input token 激活的專家數(shù)量。
• s_{i, t} 是門控機制的輸出，通過 Softmax 歸一化到 [0, 1] 區(qū)間，表示第 t 個 token 選擇第 i 個專家的概率。向量 u_t 是第 t 個 token 的輸入隱藏狀態(tài)，而 e_i 是第 i 個專家的“質(zhì)心”，可以看作歷史上路由到第 i 個專家的 token 嵌入平均值。因此，s_{i, t} 度量的是當前輸入與第 i 位專家歷史接收 token 的平均值的接近程度。
• 因此，P_i 可視為整個輸入序列選擇第 i 個專家的平均概率。
• f_i 表示被路由到第 i 個專家的 token 比例。

需要注意的是，f_i 是不可微分的，因此最小化上述損失函數(shù)實際上轉(zhuǎn)化為了最小化 s_{i, t}。同時由于 f_i 依賴于 s_{i, t}，對 s_{i, t} 的調(diào)整也會影響 f_i，從而實現(xiàn)對各專家負載分配的調(diào)節(jié)。

然而，用這種輔助損失函數(shù)來均衡負載需要付出一定代價，因為其梯度可能會干擾語言建模目標（language modeling objective）的梯度，導致模型性能下降，在極端不平衡情況下（工作負荷過大的專家的 f_i 和 P_i 都變得極大時）尤其明顯。

因此，采用這種方法進行負載均衡需要謹慎設置輔助損失函數(shù)的權(quán)重。為更清晰地說明這一點，文獻[5]的作者進行了一個實驗，用不同 alpha 值訓練模型，結(jié)果如下圖所示，其中縱軸表示困惑度指標下的模型性能，橫軸表示 MaxVio（衡量負載不平衡程度的指標，MaxVio 值越高表示負載越不平衡，i 表示第 i 個專家）：

圖 3. 輔助損失函數(shù)控制訓練中負載均衡與模型性能的權(quán)衡困境。圖片引自文獻[5]。

如圖所示，當 alpha 過小時（alpha=0），MaxVio 保持高位，說明輔助損失函數(shù)未能有效實現(xiàn)負載均衡目標。另一方面，當 alpha 過大時（alpha=0.01），模型最終會產(chǎn)生更高的困惑度。

綜上，輔助損失函數(shù)控制的負載均衡是把雙刃劍：若 alpha 未經(jīng)仔細調(diào)校，可能損害模型性能。實際 LLM 訓練中，由于資源限制，alpha 的調(diào)校過程充滿挑戰(zhàn)，這進一步增加了優(yōu)化難度。

上圖同時展示了本文提出的無損失函數(shù)方法在相同 Perplexity-MaxVio 坐標系下的表現(xiàn)，該方法同時實現(xiàn)了低困惑度和低 MaxVio，證明了無損失函數(shù)方法的有效性。

1.4 專家選擇（Expert Choice）策略

在此需提及的另一項前期工作是 Expert Choice [7]，它提出了一種簡單高效的負載均衡方法，將路由策略從“token choice”切換為“expert choice”。

具體而言，MoE 路由中的門控分數(shù)通常通過對 affinity matrix（譯者注：二維矩陣，用于量化 input token 與各個專家之間的匹配程度。）應用 Softmax 計算得出，如圖 2 所示。傳統(tǒng)路由方法從 token 維度應用 Softmax 為每個 token 選擇專家，因此這些方法被稱為“token choice”。問題在于，該機制下我們無法控制每個專家接收的 token 數(shù)量，最終導致負載不均衡問題。

而 Expert Choice 方法則從專家維度應用 Softmax，為每個專家選擇被路由的 token。通過這種設計，每個專家接收的 token 數(shù)量能夠天然達到完美均衡，因此無需依賴輔助損失函數(shù)實現(xiàn)負載均衡。在文獻[7]中，這種方法同時展現(xiàn)出更優(yōu)的模型性能和更快的訓練速度。

然而，Expert Choice 這種方法也存在局限 —— 未來 token 泄露問題。由于每個專家需要先查看所有 token 的路由分數(shù)才能決定處理哪些 token，這違反了因果關(guān)系（causality），在文本生成、機器翻譯等自回歸任務中可能引發(fā)嚴重問題。

2.DeepSeek 的無輔助損失函數(shù)的負載均衡

為在不引入 gradient inference（譯者注：此處或為作者筆誤？應當為“gradient interference（梯度干擾）”，多個損失函數(shù)（或多個優(yōu)化目標）的梯度方向發(fā)生沖突。） 的情況下解決負載均衡問題，DeepSeek 提出了一種名為" Loss-Free Balancing"的新技術(shù)，通過直接調(diào)整門控分數(shù) s_{i,t} 實現(xiàn)。

如前文所述，當我們最小化圖 2 所示的輔助損失函數(shù)時，最終會通過調(diào)整 s_{i,t} 來實現(xiàn)最小化 P_i。

因此，若能直接調(diào)整 s_{i,t}，理論上應能達到與施加輔助損失函數(shù)相似的效果。

為此，我們在每個專家的門控分數(shù)上添加了專家層面的偏置項，如下圖所示。需注意 b_i 并不用于最終門控分數(shù)的計算（后文也將說明該偏置項也是不可微分的），而是用于 TopK 選擇專家時：

圖 4. 在門控分數(shù)中引入偏置項 b_i。圖片引自文獻[5]

該偏置項 b_i 的計算方式非常直觀，如下圖所示：首先獲取分配給各專家的 token 數(shù)量平均值及所有專家的理論全局均值，然后計算給各專家分配的 token 數(shù)量與理論全局均值的差值，偏置項由該差值（或誤差）的符號乘以固定更新率（fixed update rate）決定（該更新率為可調(diào)超參數(shù)）。后續(xù)章節(jié)我們將對該超參數(shù)的影響進行更多實驗分析。

圖 5. DeepSeek 無損失函數(shù)的負載均衡算法。圖片引自文獻[5]

現(xiàn)可通過下表總結(jié)不同負載均衡方法的優(yōu)勢與局限：

圖 6. 不同負載均衡方法對比。圖片引自文獻[5]

圖 3 已展示該方法在模型性能與負載均衡間取得了更好的權(quán)衡，但仍有多方面需要驗證。下一章節(jié)我們將深入分析實驗結(jié)果。

3.Evaluation

有三個關(guān)鍵問題需要回答：

• DeepSeek 所提出的方法能否在性能和負載均衡之間實現(xiàn)更好的權(quán)衡？
• 圖 5 中更新率 u 有什么影響？
• 我們能否進一步優(yōu)化偏置項更新規(guī)則（鑒于其如此之簡單）？

3.1 性能 vs. 負載均衡

為回答第一個問題，作者在 1B 和 3B 模型上進行了實驗，比較 loss-controlled 負載均衡和 loss-free 負載均衡的困惑度（Perplexity）和 MaxVio，結(jié)果如下圖所示：

圖 7. loss-controlled 負載均衡和 loss-free 負載均衡的對比。圖片來自

上述結(jié)果與我們在圖 3 中看到的結(jié)果類似：所提方法同時實現(xiàn)了更低的困惑度和更低的MaxVio。

除了評估最終的 checkpoint 外，作者還展示了訓練過程中的 MaxVio 曲線，以便更全面地理解該方法在整個訓練過程中的表現(xiàn)，如下圖所示：

圖 8. 訓練過程中的 MaxVio 曲線。圖片來自文獻[5]。

如圖中所示，在 1B 和 3B 的模型配置下，loss-free 方法在整個訓練過程中都展現(xiàn)出更優(yōu)的負載均衡能力，體現(xiàn)了該方法的穩(wěn)定性。

3.2 超參數(shù)的影響（更新率）

如圖 5 所示，所提方法引入了一個新的超參數(shù) u（稱為更新率（update rate）），該超參數(shù)如何影響 loss-free 方法的有效性？具體而言，我們需要理解 loss-free 方法對超參數(shù) u 的取值是敏感還是不敏感，以及如何選擇一個最優(yōu)值來最大化該方法的效果。

如前文所述，在門控分數(shù)中添加偏置項的概念類似于繞過損失函數(shù)的反向傳播，直接對門控分數(shù)進行梯度更新。在這種情況下，更新率 u 的作用與梯度更新中的步長（step size）類似。由此可推測其影響也相似：過小的更新率可能會導致收斂速度緩慢。過大的更新率可能導致不穩(wěn)定和引發(fā)波動。

在原文中，作者對更新率進行了實驗測試（取值從 1e-4 到 1e-2），結(jié)果如下圖所示：

圖 9. 更新率（update rate）對訓練負載均衡的影響。圖片來自文獻[5]。

與預期一致，當 u 過小時（如 1e-4），MaxVio 下降速度較慢；而過大的 u（如 1e-2）則因波動性增強，導致訓練過程中 MaxVio 持續(xù)偏高。

3.3 其他偏置項更新規(guī)則

為回答第三個問題，研究者嘗試了多種備選策略，并將它們與 DeepSeek 提出的版本進行對比：

• 策略變體 1：使用 e_i 的數(shù)值（而不僅僅是符號）計算偏置項，即從 b_i = b_i +u?sign(e_i) 改為b_i = b_i +u?e_i。
• 策略變體 2：使用乘法偏置項而非加法偏置項。

其中策略變體 2 可以更正式地描述如下：

image.png

實驗表明，策略變體 1 能帶來略優(yōu)的負載均衡效果，但未提升模型性能：

圖 10. 策略變體 1 的性能表現(xiàn)。圖片來自文獻[5]。

而策略變體 2 甚至顯示出略差的模型性能：

圖 11. 策略變體 2 的性能表現(xiàn)。圖片來自文獻[5]。

以上所有結(jié)果均表明，最簡單的策略反而是最佳選擇。

4.Summary

在本文中，我們解釋了 DeepSeek 模型的核心架構(gòu)創(chuàng)新之一 —— DeepSeekMoE 中使用的無輔助損失函數(shù)的負載均衡方法。

本文首先介紹了混合專家模型（MoE）的基本原理，強調(diào)了負載均衡的重要性，并回顧了先前的解決方案（包括 auxiliary loss 方法和 Expert Choice 機制）。接著，本文闡釋了 DeepSeek 的無損失函數(shù)的負載均衡方法及其性能表現(xiàn)。

DeepSeek 的無損失函數(shù)方法在保持因果關(guān)系的同時避免了引入梯度干擾，其有效性已通過原論文的實證結(jié)果得到驗證。

感謝您花時間閱讀本文！

參考文獻

[1] DeepSeek（https://www.deepseek.com/）

[2] DeepSeek-V3 Technical Report（https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf）

[3] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model（https://arxiv.org/abs/2405.04434）

[4] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models（https://arxiv.org/abs/2401.06066）

[5] Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts（https://arxiv.org/abs/2408.15664）

[6] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding（https://arxiv.org/abs/2006.16668）

[7] Mixture-of-Experts with Expert Choice Routing（https://arxiv.org/abs/2202.09368）

Thanks for reading! 

Hope you have enjoyed and learned new things from this blog!