Technical Summary: DOT-MoE: Differentiable Optimal Transport for MoEfication

1. Problem Statement

The rapid scaling of Large Language Models (LLMs) has created a significant bottleneck in inference efficiency. While Mixture-of-Experts (MoE) architectures decouple model size from inference cost by activating only a subset of parameters per token, training MoEs from scratch is notoriously unstable, data-hungry, and computationally expensive.

Consequently, MoEfication—the conversion of pre-trained dense models into sparse MoEs—has emerged as a viable alternative. However, existing methods for converting dense Feed-Forward Networks (FFNs) into sparse experts rely heavily on heuristic strategies (e.g., random splitting, weight-based clustering, or activation-based clustering). These approaches suffer from three primary limitations:

1. Separation of Concerns: They typically treat neuron assignment (partitioning the FFN) and router training as separate, sequential stages, preventing the two from co-adapting.
2. Proxy Optimization: They optimize based on intermediate statistics (weights or activations) rather than the actual output contribution of neurons, often leading to suboptimal reconstruction of the dense model's behavior.
3. Lack of Differentiability: The discrete nature of neuron assignment prevents end-to-end gradient flow, making it difficult to jointly optimize the structural decomposition and the routing policy.

2. Methodology: DOT-MoE

The authors propose DOT-MoE, a framework that formulates the decomposition of dense layers into experts as a Differentiable Optimal Transport (DOT) problem. This approach enables end-to-end learning of both the expert decomposition and the routing mechanism.

2.1. Optimal Transport Formulation

The core insight is to view neuron assignment as a balanced transport problem where "mass" (neurons) must be transported from a source distribution (the dense FFN neurons) to a target distribution (the experts) under strict capacity constraints.

• Source: $dffn$ neurons, each with unit mass.
• Target: $E$ experts, each requiring exactly $s$ neurons ($s \cdot E = dffn$).
• Cost Function: Instead of a fixed cost matrix, the method uses a learnable affinity matrix $A$, where $A_{i,e}$ represents the affinity of assigning neuron $i$ to expert $e$. The objective is to maximize the total affinity while satisfying marginal constraints.

2.2. Differentiable Relaxation via Sinkhorn-Knopp

To make the discrete assignment problem differentiable, DOT-MoE employs Entropic Regularization.

• The optimal transport problem is relaxed using the Sinkhorn-Knopp algorithm, which iteratively normalizes rows and columns of the transport plan matrix.
• This yields a soft assignment matrix $M_{soft} \in [0, 1]^{dffn \times E}$ that satisfies capacity constraints in expectation.
• As the temperature parameter $\tau \to 0$, the solution approaches the discrete optimum.

2.3. Joint Optimization with Straight-Through Estimators (STE)

A critical challenge is that deployment requires discrete assignments, and routing involves discrete top-k selection, both of which are non-differentiable. DOT-MoE addresses this using Straight-Through Estimators (STE):

• Forward Pass: Uses hard decisions (greedy rounding of $M_{soft}$ for assignment; top-k selection for routing).
• Backward Pass: Routes gradients through the soft counterparts ($M_{soft}$ and the router logits $P$).
• This allows the router weights ($W_r$) and the assignment affinity matrix ($A$) to co-adapt during training.

2.4. Training Objective

The framework is trained via an Alignment Phase where the dense model weights are frozen. The model minimizes a combination of losses:

1. Output Reconstruction: Measured via KL divergence between the dense teacher and the sparse MoE student outputs.
2. Language Modeling: Cross-entropy loss on the language modeling objective.
3. Stability Auxiliaries: Router $z$-loss (to prevent logit explosion) and load-balancing loss (to prevent expert collapse).

3. Key Contributions

1. Optimal Transport Framework: The paper introduces a novel formulation of dense-to-MoE conversion as a balanced transport problem, solved via differentiable Sinkhorn iterations to enforce strict expert capacity constraints.
2. Dual-Level Co-Adaptation: Unlike prior methods that fix assignments before training routers, DOT-MoE jointly optimizes neuron-to-expert decomposition and token-to-expert routing using complementary STEs.
3. Output-Aware Construction: By directly modeling the neuron's contribution to the final output (rather than optimizing proxies like weights or intermediate activations), the method achieves higher reconstruction fidelity.

4. Experimental Results

The authors evaluated DOT-MoE across three model families (LLaMA-2, LLaMA-3, Qwen2.5) and six benchmarks (ARC-Challenge, Winogrande, HellaSwag, PIQA, SciQ, BoolQ).

• Performance vs. Pruning: On LLaMA-2 7B with a 50% parameter budget, DOT-MoE achieved a perplexity of 7.99 on WikiText-2, significantly outperforming the state-of-the-art structured pruning method DISP-LLM (9.84).
• Performance vs. Existing MoEfication: In zero-shot settings, DOT-MoE substantially outperformed methods like CMoE and LLaMA-MoE-v2. For instance, on LLaMA-2 7B, DOT-MoE achieved 61.5% average accuracy compared to CMoE's 44.5%.
• Efficiency: DOT-MoE retains approximately 90% of the original dense model's performance while reducing active parameters by 50%.
• Fine-Tuning: With only 1.2B tokens of continuous fine-tuning, DOT-MoE models recovered performance close to dense counterparts (e.g., 66.6% vs 76.6% on LLaMA-2 7B), outperforming baselines that required similar or larger data budgets.
• Ablation Studies:
  • Granularity: Performance improves with increased expert granularity up to a saturation point, after which gains diminish.
  • Inference Speed: Despite higher expert counts, inference throughput remains stable when using fused MoE kernels (e.g., vLLM), as the total fused intermediate dimension remains constant.
  • Sparsity: Models trained at higher sparsity levels (75%) generalize robustly to even more extreme sparsity regimes at inference time.

5. Significance and Claims

The paper claims that DOT-MoE represents a significant advancement in the efficiency of LLM inference by providing a principled, differentiable method for converting dense models into sparse MoEs.

• Paradigm Shift: It moves away from heuristic, static partitioning toward a dynamic, learnable decomposition that respects the functional coherence of experts.
• Practicality: The method enables efficient post-hoc sparsification of pre-trained models without the need for full retraining from scratch, retaining high capacity while drastically reducing active parameters.
• Robustness: The framework demonstrates that preserving FFN fidelity requires explicitly modeling the neuron's contribution to the output, a principle that existing proxy-based methods fail to capture.

The authors conclude that DOT-MoE offers a superior initialization for sparse models, bridging the performance gap between dense and sparse architectures with minimal additional training data.