DOT-MoE: MoE化のための微分可能最適輸送

要旨

大規模言語モデル（LLM）のスケーリングは大幅な性能向上をもたらしたが、推論効率において重大な課題を生み出している。混合専門家モデル（MoE）アーキテクチャは、モデルサイズと推論コストを切り離すことでこの問題に対処するが、MoEをゼロから学習することはしばしば不安定であり、計算負荷が大きい。事前学習済みの高密度モデルを疎なMoEに変換することが代替解決策として浮上しているが、既存手法は通常、フィードフォワードネットワーク（FFN）を専門家に分割するためにヒューリスティックなニューロンクラスタリングやランダム分割に依存している。本研究では、高密度層の分解を微分可能最適輸送（DOT）問題として定式化する新しいフレームワークDOT-MoEを提案する。静的ヒューリスティックの代わりに、ニューロン割り当てをバランスのとれた輸送問題としてモデル化し、微分可能なSinkhorn-Knopp反復を利用して厳密な専門家容量制約を強制する。さらに、ストレートスルー推定器（STE）を活用し、離散的なニューロン対専門家の割り当てと、トークン対専門家のルーティングポリシーをエンドツーエンドで共同学習する。複数のアーキテクチャとベンチマークにわたる広範な実験により、DOT-MoEは構造化プルーニング、ヒューリスティッククラスタリング、ランダム分割のベースラインを大幅に上回り、元の高密度モデルの性能の90%を維持しつつ、アクティブパラメータを50%削減することを示す。

English

The scaling of Large Language Models (LLMs) has driven significant performance gains but created substantial challenges in inference efficiency. While Mixture of Experts (MoEs) architectures address this by decoupling model size from inference cost, training MoEs from scratch is often unstable and compute intensive. Conversion of pre-trained dense models into sparse MoEs has emerged as an alternative solution; however, existing methods typically rely on heuristic neuron clustering or random splitting to partition the Feed-Forward Network (FFN) into experts. In this work, we propose DOT-MoE, a novel framework that formulates the decomposition of dense layers as a Differentiable Optimal Transport (DOT) problem. Instead of static heuristics, we model neuron assignment as a balanced transport problem, utilizing differentiable Sinkhorn-Knopp iterations to enforce strict expert capacity constraints. Furthermore, we utilize Straight-Through Estimators (STE) to jointly learn the discrete neuron-to-expert assignment and the token-to-expert routing policy end-to-end. Extensive experiments across multiple architectures and benchmarks demonstrate that DOT-MoE significantly outperforms structured pruning, heuristic clustering, and random-split baselines, retaining 90% of the original dense model's performance while reducing active parameters by 50%.