[論文レビュー] Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

Lean Wang, Huazuo Gao+3 more|arXiv (Cornell University)|Aug 28, 2024

Distributed Sensor Networks and Detection Algorithms被引用数 6

ひとこと要約

本論文は、Loss-Free Balancing（ロスフリーバランシング）を提案します。これは、補助損失を用いずにMoEの負荷バランスを制御する手法で、専門家ルーティングスコアを動的にバイアスすることにより、パープレキシティと負荷バランスを改善します。

ABSTRACT

For Mixture-of-Experts (MoE) models, an unbalanced expert load will lead to routing collapse or increased computational overhead. Existing methods commonly employ an auxiliary loss to encourage load balance, but a large auxiliary loss will introduce non-negligible interference gradients into training and thus impair the model performance. In order to control load balance while not producing undesired gradients during training, we propose Loss-Free Balancing, featured by an auxiliary-loss-free load balancing strategy. To be specific, before the top-K routing decision, Loss-Free Balancing will first apply an expert-wise bias to the routing scores of each expert. By dynamically updating the bias of each expert according to its recent load, Loss-Free Balancing can consistently maintain a balanced distribution of expert load. In addition, since Loss-Free Balancing does not produce any interference gradients, it also elevates the upper bound of model performance gained from MoE training. We validate the performance of Loss-Free Balancing on MoE models with up to 3B parameters trained on up to 200B tokens. Experimental results show that Loss-Free Balancing achieves both better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies.

研究の動機と目的

• ルーティング崩壊と過剰な計算を避けるために、MoEにおける専門家ルーティングのバランスの必要性を動機づける。
• 干渉勾配を導入することなく、負荷バランスを制御する Loss-Free Balancing を提案する。
• 本法が1Bおよび3BのMoEモデルにおいて、補助損失ベースの手法よりも優れた検証時のパープレキシティと負荷バランスをもたらすことを実証する。
• 専門家並列性との互換性を示し、代替手法と比較した理論的リークを分析する。

提案手法

• Top-Kルーティングの前段でゲーティングスコアに専門家ごとのバイアスを導入する。
• 観測された専門家の負荷に基づいて各バイアス b_i を反復的に更新し、過負荷の専門家を抑制し、低負荷の専門家を高める。
• 訓練バッチごとに勾配を主目的に加えず、アルゴリズム1を用いてバイアスを更新する。
• ベンチマークMoE設定で、Loss-Free Balancing を補助損失で制御された負荷バランシングおよび Expert Choice と比較する。
• 検証時のパープレキシティとグローバルな負荷バランスメトリック MaxVio（グローバルおよびバッチ版）を用いて評価する。
• 専門家並列性との互換性を議論し、将来トークンのリーク懸念を分析する。

実験結果

リサーチクエスチョン

• RQ1 補助損失なしのゲーティングスコアのバイアス付けは、補助損失よりも良い負荷バランスを生み出すか？
• RQ2 Loss-Free Balancing は1Bおよび3BのMoEモデルで、モデル性能（パープレキシティ）と負荷バランスの両方を改善できるか？
• RQ3 この手法は専門家並列性と互換性があり、干渉勾配を生じないか？
• RQ4 将来トークンのリークとトレーニングの安全性の観点で、Loss-Free Balancing は Expert Choice とどう比較されるか？

主な発見

モデルサイズ	負荷バランシング手法	検証パープレキシティ	MaxVio_global
1B	Loss-Controlled	9.56	0.72
1B	Loss-Free	9.50	0.04
3B	Loss-Controlled	7.97	0.52
3B	Loss-Free	7.92	0.04

• Loss-Free Balancing は1Bおよび3Bモデルのいずれにおいても、補助損失制御手法より低い検証時パープレキシティを達成する。
• グローバル MaxVio 負荷バランスメトリックは、Loss-Free Balancing では0.04と大幅に低く、補助損失では1Bで0.72、3Bで0.52である。
• Loss-Free BalancingはMaxVio_batch曲線で一貫した優位性を示し、訓練全体を通じて負荷バランスをより良く維持する。
• この方法は専門家並列性と互換性があり、計算・バッチサイズが大きくなるにつれてバランスが改善されることを示す。
• 加法的な専門家バイアスは、乗法的なバイアスよりもバランスと性能で優れる。
• Softmaxゲート実験では、その設定で Loss-Free Balancing がより良いバランスとわずかに良いパープレキシティを示す。