Udbhav Bamba, Arnav Chavan, Aryamaan Thakur, Steve Teig, Deepak Gupta
7 min
Abstract
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%.
Alex: So the math essentially melts the hard choices into something fluid, so the model can slide toward the best configuration rather than having to guess it perfectly from the start.
Sam: That's a good way to put it. And there's one more piece: something called a "Straight-Through Estimator." During training, the model needs to make real, committed decisions — a neuron does end up in one specific group. But the Straight-Through Estimator lets the model behave as if those decisions were still soft and adjustable when it's calculating how to improve. It's a useful trick that lets the groupings and the routing system evolve together, rather than one being locked in before the other is ready.
Alex: So the groups and the system that directs traffic between them are being refined at the same time. Why does that matter so much compared to older methods?
Sam: Because older methods were optimizing for the wrong thing. They looked at a neuron's internal settings — its weights — or how often it fired, and used those as a stand-in for how useful the neuron actually was. But that's a bit like picking a football team based on how fast players run in practice, rather than how well they perform in actual matches. The DOT-MoE approach looks at a neuron's real contribution to the model's output — what it actually adds to the final answer — and that turns out to be a much more meaningful signal.
Alex: So previous methods were measuring potential, while this one measures actual impact.
Sam: Exactly. And that shift in what you measure makes a meaningful difference in the quality of the groups you end up with. The paper reports that the resulting system retains around ninety percent of the original model's performance while cutting the number of active neurons in half for any given task.
Alex: That's a notable result. Does the model get there quickly, or does it take a long time to settle into a good configuration?
Sam: The paper tracks this directly. Their system shows lower error rates earlier in training compared to competing methods — it starts well and stays ahead. Some other approaches do eventually catch up on training data, but they show signs of what's called "overfitting." Think of a student who memorizes every answer on a practice exam but hasn't actually understood the material — they do well on the test they studied, but struggle with anything new. The experts in DOT-MoE, because they're shaped by real output signals, tend to learn more general, transferable patterns rather than just memorizing specifics.
Alex: What about sparsity itself — if you force the model to use fewer experts at once, does that make learning harder?
Sam: Interestingly, the paper suggests the opposite. Training under higher sparsity — where the model is constrained to use fewer resources — tends to produce more robust representations. It's a bit like being asked to summarize a long chapter in a single paragraph. The constraint forces you to identify what actually matters and discard the rest. The model, under that pressure, learns to encode information more efficiently.
Alex: So the constraint itself becomes a useful discipline.
Sam: That seems to be the finding, yes. Though the authors are careful to note limitations. The optimization process starts from a random initial state, which means it takes time to converge. They suggest that using data-informed starting points — rather than random ones — could speed things up considerably. That's flagged as a direction for future work.
Alex: And are there other ways to push efficiency further beyond what this paper covers?
Sam: One avenue they mention is what they call "hard pruning" — permanently deleting experts that rarely get used. Right now, all the expert groups stay loaded in memory even if only a few are active at any moment. Actually removing the unused ones would reduce the memory footprint of the model itself, not just the computation. That's a separate problem from what this paper solves, but it's a natural next step.
Alex: So the paper establishes a more principled way to convert dense models into sparse ones, and points toward a broader research agenda for making these systems leaner without sacrificing capability.
Sam: That's a fair summary. The core contribution is replacing rigid, rule-based grouping with a learned, output-aware process — and showing that the difference in quality is meaningful. The open questions around initialization and pruning suggest there's still room to go further.
Alex: Thanks for walking through that, Sam. And thanks to everyone listening to ResearchPod.