DOT-MoE: Differentiable Optimal Transport for MoEfication

Large language models (LLMs) are becoming increasingly expensive to serve. As model dimensions grow, dense activation patterns result in unsustainable inference latency. Standard Transformers activate every parameter for every input token. To solve this, engineers turn to Mixture-of-Experts (MoE) architectures. These decouple total model size from the actual cost of inference. They route tokens to only a small subset of specialized sub-networks.

However, training these MoE models from scratch is data-hungry and unstable. A growing alternative is "MoEfication." This involves converting a pre-trained dense model's Feed-Forward Network (FFN) layers into sparse experts. This reorganizes internal parts into specialized groups. It allows the model to run faster without the massive cost of fresh pre-training.

The difficulty lies in the "assignment problem." This means deciding which neurons belong to which expert. Most current methods use static heuristics. They might cluster neurons based on weight similarity or activation patterns. These methods often treat neuron assignment and token routing as disconnected stages. This paper argues that this decoupling is fundamentally flawed. It ignores the final output. This leads to poor reconstruction of the original model's logic.

The Problem

The status quo for dense-to-MoE conversion relies on "proxies." These are intermediate signals acting as stand-ins for how a neuron contributes to the output. Common strategies include weight-based clustering or activation-based clustering.

The authors demonstrate that these proxies are inadequate. In a controlled single-layer reconstruction analysis, they measured error rates. They found that methods relying on input-side statistics or intermediate activations incurred mean squared errors (MSE)—a measure of prediction error—ranging from 2× to over 41× higher than their proposed approach. The output of an FFN depends on the interaction between intermediate activations and down-projection weights. Optimizing for anything else misses the mark. Current methods attempt to rebuild a machine by looking at gear movement. They fail to measure the actual torque at the drive shaft.

How It Works

DOT-MoE replaces these heuristics with Differentiable Optimal Transport (DOT). The core idea treats neuron assignment as a balanced transport problem. Every neuron carries a unit of "mass" to be delivered to one expert. Every expert must receive a fixed amount of mass to meet capacity constraints.

The implementation follows three technical pillars:

1. Differentiable Sinkhorn Iterations: Ideal assignment is a discrete "hard" decision. This is non-differentiable. The authors use entropic regularization to create a "soft" assignment. They apply the Sinkhorn-Knopp algorithm—an iterative process for normalizing matrices—to find a differentiable solution. They implement this in the log-domain to maintain numerical stability.
2. Dual-Level STE Optimization: They use Straight-Through Estimators (STE) to bridge the gap between soft math and hard deployment. STE allows them to use hard, discrete assignments during the forward pass. This includes greedy rounding for neurons and top-k selection for tokens. Gradients flow through the soft, differentiable counterparts during the backward pass.
3. Output-Aware Alignment: Previous works freeze the assignment and then train a router. DOT-MoE allows the router and expert structure to co-adapt. The objective function directly measures deviation from the original dense FFN output. This ensures the sparse structure preserves the original model's residual stream (the main information path).

As shown in, this process forces experts to become functionally distinct.

The t-SNE visualization (a dimensionality reduction technique) shows experts specializing in separated regions.

Numbers

The headline result is significant. The authors report that DOT-MoE can reduce active parameters by 50%. It does this while retaining 90% of the original dense model's performance.

On the LLaMA-2 7B benchmark, DOT-MoE achieves a WikiText-2 perplexity (a measure of prediction accuracy) of 7.99 at a 50% parametric budget. This outperforms the state-of-the-art structured pruning method DISP-LLM, which scores 9.84 [Table 1]. In terms of downstream utility, the authors report a zero-shot average accuracy of 61.5% on LLaMA-2 7B. This beats the 44.5% achieved by CMoE.

Regarding infrastructure, increasing expert granularity does not necessarily hurt speed. Using vLLM’s fused MoE kernels, they show that inference throughput remains stable. This holds true even as expert counts increase, provided active parameters stay constant [Figure 1c]. This happens because fused implementations batch expert computations into large GEMMs (General Matrix Multiplications). This avoids executing many tiny, inefficient kernels.

What's Missing

While the results are strong, some gaps remain for production engineers:

• The Compute Cost of Alignment: The alignment phase is not free. The authors report a ~15% overhead per training step compared to standard dense training. Most of this comes from hard-assignment matrix construction. This is a one-time cost during conversion.
• Scale Limits: The paper evaluates scaling up to 32B parameters (Qwen2.5-32B). However, the authors did not explore the scaling behavior in the ~100B token fine-tuning regime. This was due to compute constraints.
• Hardware Bottlenecks: The current implementation of greedy rounding runs on the CPU. This creates CPU-to-GPU transfer overhead. A dedicated GPU kernel could remove much of this cost.

Should You Prototype This

Yes, if you want to increase efficiency in existing dense checkpoints.

If your team uses structured pruning or basic MoEfication, DOT-MoE offers a principled alternative. It allows you to maintain stable throughput with high expert granularity. This is a major win for users of fused kernels. However, do not expect an instant win. Budget for the alignment phase. Note that the current reliance on CPU-based rounding might slow down the initial conversion.

Figures from the paper