DOT-MoE: Differentiable Optimal Transport for MoEfication

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 Feed-Forward Networks (FFNs) 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%.

Large language models are powerful, but they are also expensive to run because every part of the model is used for every request. One way to make them more efficient is to organize them into smaller specialized components, and only activate a few of them at a time. However, building these efficient systems usually requires training entirely new models which is costly. In this work, we introduce DOT-MoE, a method for turning existing language models into more efficient sparse models without retraining them from scratch. Instead of grouping model components using simple rules or random assignments, our method learns how to organize them based on how they contribute to the model’s final predictions. By making large language models cheaper to run, our work can help lower energy usage and make advanced AI systems more accessible.