DOT-MoE: Differenzierbarer optimaler Transport für die MoEfizierung

Zusammenfassung

Die Skalierung großer Sprachmodelle (Large Language Models, LLMs) hat zwar erhebliche Leistungssteigerungen ermöglicht, aber auch beträchtliche Herausforderungen bei der Inferenzeffizienz geschaffen. Während Mixture-of-Experts-Architekturen (MoE) dieses Problem durch die Entkopplung von Modellgröße und Inferenzkosten angehen, ist das Training von MoEs von Grund auf oft instabil und rechenintensiv. Die Umwandlung vortrainierter dichter Modelle in spärliche MoEs hat sich als alternative Lösung etabliert; allerdings basieren bestehende Methoden typischerweise auf heuristischem Neuron-Clustering oder zufälliger Aufteilung, um das Feed-Forward-Netzwerk (FFN) in Experten zu partitionieren. In dieser Arbeit schlagen wir DOT-MoE vor, ein neuartiges Framework, das die Zerlegung dichter Schichten als ein differentielles optimales Transportproblem (Differentiable Optimal Transport, DOT) formuliert. Anstelle statischer Heuristiken modellieren wir die Neuronenzuweisung als ein balanciertes Transportproblem und nutzen differenzierbare Sinkhorn-Knopp-Iterationen, um strikte Expertenkapazitätsbeschränkungen durchzusetzen. Darüber hinaus verwenden wir Straight-Through-Schätzer (STE), um die diskrete Neuron-zu-Experte-Zuweisung und die Token-zu-Experte-Routing-Politik end-to-end gemeinsam zu erlernen. Umfangreiche Experimente über mehrere Architekturen und Benchmarks hinweg zeigen, dass DOT-MoE strukturiertes Pruning, heuristisches Clustering und zufällige Aufteilungen als Baseline signifikant übertrifft, wobei 90% der ursprünglichen Leistung des dichten Modells erhalten bleiben, während die aktiven Parameter um 50% reduziert werden.

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%.