MoEは負荷分散だけでは不十分：DRMoETがルーティングの頑健性を強化

導入

Mixture-of-Experts（MoE）トランスフォーマーは、各トークンに対して一部のエキスパートだけを起動することで、計算量を大きく増やさずにモデル容量を拡張します。一方で、疎なルーティングには見えにくい信頼性の問題があります。ルーターがトークンを均等に振り分けても、各エキスパートが割り当てられた入力を十分に学習できているとは限らないためです。

論文「Distributionally Robust Mixture-of-Experts Training」は、この問題に対してDRMoETを提案しています。狙いは、単純な利用率の均等化から一歩進み、ルーティングが完全でない場合にも性能を保てるようにすることです。

方法：高損失の経路を学習する

従来の負荷分散は、エキスパートごとのトークン数を調整し、特定のエキスパートだけが過負荷になったり、他がほとんど使われなかったりする状態を避けます。しかし、これはエキスパートが担当データをうまく処理できているかを直接評価しません。

DRMoETは各層で、実際に活性化されたエキスパートの損失を追跡します。入力による重み付けを行った損失を指数移動平均（EMA）で平滑化し、エントロピー正則化付きのSoftmaxルールによってエキスパート分布を更新します。損失が高いルーティング結果には、より大きな学習上の注意が向けられます。

ここで重要なのは、推論時により多くのエキスパートを起動する方法ではないことです。通常のTop-kによる疎な計算は維持されます。学習時に、現在の最上位候補だけでなく、実際に選ばれる可能性がある非Topの経路も十分に訓練することで、ルーターの不確実性に対する余裕を作ります。

実験結果

著者らはFLAME-MoEの訓練レシピを用い、総パラメータ数746Mと10.3Bの2つの規模で評価しました。DRMoETは、標準FLAME-MoEと補助損失を使わない負荷分散方式の双方に対して、下流タスクの平均値を改善したと報告されています。

10.3Bモデルを67Bトークンで学習した設定では、7タスク平均が標準条件の0.6625からDRMoETで0.6767に上昇しました。補助損失なしのベースラインは0.6431でした。分析では、平均損失をほぼ変えずにエキスパート損失の分散を下げたことが示されています。また、中位順位のエキスパートを強制的に使う誤ルーティング試験では、余分な損失が4.3%低下し、ドメインごとの専門化も改善しました。

意義と限界

この研究は、MoEの評価軸に新しい視点を加えます。トークンの流量が均衡していても、エキスパート間の能力差が大きければ、ルーターの小さな判断ミスが性能低下につながります。利用率の均衡に加えて、候補エキスパートがどれだけ入力をカバーできるかを最適化することが、疎なモデルの信頼性向上に役立つ可能性があります。

ただし、報告された結果は特定のレシピ、規模、タスクに基づくものです。別のルーター構造やデータ分布、さらに大規模な学習でも効果が再現するかは今後の検証課題です。エントロピー正則化や損失平滑化の設定も、実装時の安定性に影響すると考えられます。

出典：Hugging Face Daily Papers