PaperlayerView exact revision on arXiv

Source-linked AI summary

DOT-MoE: Differentiable Optimal Transport for MoEfication

Udbhav Bamba, Arnav Chavan, Aryamaan Thakur, Steve Teig, Deepak Gupta

arXiv:2606.01666v1cs.LGcs.AI

TL;DR

Converting dense models into sparse MoEs is limited by heuristic partitioning and the coupled problem of assigning neurons and routing tokens. DOT-MoE uses differentiable optimal transport to jointly learn both, outperforming pruning and existing MoEfication methods while retaining 90% of dense-model performance at 50% parametric count.

  • Problem

    Existing dense-to-MoE methods use heuristic or proxy-based neuron assignments, despite the tight coupling between expert decomposition and token routing.

  • Method

    DOT-MoE formulates neuron-to-expert assignment as balanced differentiable optimal transport and jointly learns assignments with token routing using Sinkhorn iterations and straight-through estimators.

  • Results

    90% of the original dense model’s performance is retained at 50% parametric count, with DOT-MoE outperforming structured pruning and existing MoEfication methods across three model families and six benchmarks.

  • Takeaways & Limitations

    DOT-MoE achieves a superior sparsity–performance trade-off compared with heuristic clustering and structured pruning methods.

  • Takeaways & Limitations

    Existing assignment methods optimize intermediate-representation proxies while neglecting how neuron activations and down-projection weights interact to produce the final output.

Abstract

from arXiv · show

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

1. Introduction

LLM scaling makes dense inference increasingly costly, while training MoEs from scratch is unstable and data-hungry. DOT-MoE addresses dense-to-MoE conversion by jointly learning balanced neuron assignment and token routing through differentiable optimal transport.

  • Motivation: Dense Transformers activate every parameter for every token, creating unsustainable inference latency and resource consumption as model dimensions grow.MoEs reduce inference cost by routing tokens to a small subset of expert sub-networks.
  • Motivation: Training MoEs from scratch is data-hungry and unstable, motivating conversion of pre-trained dense checkpoints into sparse MoEs.MoEfication transforms the Feed-Forward Network in each block while leveraging existing dense-model representations.
  • Challenge: Neuron assignment is the core conversion challenge, but existing methods largely use heuristics and separate assignment from router training.This separation prevents a unified, differentiable treatment of expert structure and routing.
  • DOT-MoE: DOT-MoE formulates dense-layer decomposition as a Differentiable Optimal Transport problem and learns expert decomposition and routing end-to-end with differentiable Sinkhorn-Knopp iterations.Unlike prior methods that freeze assignment before router training, DOT-MoE jointly optimizes both components.
  • DOT-MoE: The framework uses balanced transport and complementary straight-through estimators to jointly optimize neuron-to-expert decomposition and token-to-expert routing.This dual-level mechanism enables the router and expert structure to co-adapt.

2. Background and Motivation

Dense-to-MoE conversion reduces per-token computation by partitioning FFN neurons into sparse expert groups, but existing assignment strategies optimize input-side or intermediate statistics rather than FFN output fidelity. A controlled reconstruction analysis shows these strategies can produce substantially higher errors than the proposed approach.

  • Background and Motivation: FFNs typically account for approximately two-thirds of transformer language-model parameters, making them a primary target for efficiency improvements.The FFN uses gate, up-projection, activation, element-wise multiplication, and down-projection operations.
  • Background and Motivation: MoEfication partitions d_ffn intermediate neurons into E disjoint experts of size s = d_ffn/E, with a learned router selecting the top-k experts per token.This preserves total model capacity while reducing per-token computation from O(d_ffn) to O(k · d_ffn/E).
  • Background and Motivation: Existing methods assign neurons randomly, by projection-weight similarity, or through activation, gradient, and co-activation statistics.Random assignment requires extensive continued pre-training, while clustering methods rely on input-side weights or intermediate activation patterns.
  • Background and Motivation: Existing approaches neglect the interaction between intermediate activations H and down-projection weights W_down that determines FFN output.Structured pruning uses layer-wise importance scores, while dense-to-MoE methods cluster input weights, activations, or co-activation patterns.
  • Background and Motivation: 2× to over 41× higher mean squared errors were observed for input-side and intermediate-activation methods than for the proposed approach in single-layer reconstruction tests.The analysis covered LLaMA-2 and LLaMA-3 and isolated expert construction strategy.

3. Method

DOT-MoE converts pretrained dense FFN layers into balanced sparse MoE layers by jointly learning neuron assignments and token routing through differentiable optimal transport. Entropic Sinkhorn relaxation, hard capacity-constrained assignments, and straight-through estimators enable end-to-end optimization without full model fine-tuning.

  • Method: DOT-MoE partitions dffn neurons into E experts of s neurons each, activates k < E experts per token, and targets post-hoc sparsification without full fine-tuning.The balance condition is s · E = dffn, with only k experts active for each token.
  • Optimal Transport Assignment: Neuron assignment is formulated as balanced optimal transport, assigning each neuron once and requiring every expert to receive exactly s neurons.A learnable affinity matrix A defines assignment costs as C = −A, maximizing total affinity under capacity constraints.
  • Differentiable Relaxation: Standard optimal transport is non-differentiable and computationally prohibitive, so entropic regularization produces a differentiable interior solution computed by log-domain Sinkhorn-Knopp iterations.The regularized solution approaches the unregularized optimum as τ → 0 and the uniform plan as τ → ∞.
  • Token Routing: A linear router applies top-k selection to choose the experts activated for each input token.Router weights Wr parameterize token-to-expert routing, while the selected expert indices form the routing mask.
  • Joint Optimization: Greedy rounding creates disjoint binary expert clusters with |Ce| = s, while straight-through estimators use hard forward decisions and soft backward gradients for assignment and routing.Reconstruction-loss gradients update the affinity matrix A through Msoft and router weights Wr through P.
  • Training and Conversion: Training masks intermediate activations so only k · s neurons contribute per token and jointly optimizes assignment and routing with teacher KL, language-model cross-entropy, and stability losses.After convergence, the final assignment is extracted into a standard sparse MoE; the balanced-transport formulation also extends to grouping attention heads.

4. Experiments

Across pruning and dense-to-MoE comparisons, DOT-MoE achieves stronger performance retention and favorable scaling while preserving inference flexibility and throughput. Its output-aware initialization also improves training generalization relative to heuristic conversions.

  • Comparison with Pruning Methods: 7.99 perplexity is the lowest among evaluated methods at a 50% parametric budget, versus 9.84 for DISP-LLM.This comparison is reported on LLaMA-2 7B.
  • Comparison with Pruning Methods: 72.3% average accuracy versus 66.7% for DISP-LLM shows stronger knowledge retention on Qwen2.5-7B.DOT-MoE outperforms pruning baselines across all three model families.
  • Comparison with Dense-to-MoE Methods: 61.5% average accuracy on LLaMA-2 7B exceeds CMoE’s 44.5% through jointly learned assignment and routing.This enables stronger zero-shot transfer without training model weights during alignment.
  • Comparison with Dense-to-MoE Methods: 71.0% accuracy on LLaMA-3 8B with 7B tokens exceeds 66.8% for LLaMA-MoE-v2 trained on a similar data split.On LLaMA-2 7B, scaling tokens increases accuracy from 61.5% to 66.6%.
  • Scaling Results: +34.3 points over CMoE on Qwen2.5-32B demonstrates scaling to larger models, with gains maintained across context lengths up to 32K tokens.The reported pre-trained model average accuracy is 80.6%.
  • Inference Efficiency: No throughput penalty occurs as expert count increases when active parameters remain constant and fused MoE kernels are used.Throughput remains stable because expert computations are batched into large GEMMs.

5. Future Work and Conclusion

DOT-MoE converts dense FFNs into sparse MoEs through differentiable optimal transport, jointly learning neuron-to-expert assignments and routing policies. Future work targets data-driven initialization and pruning persistently underutilized experts.

  • Conclusion: DOT-MoE formulates dense-FFN conversion as a Differentiable Optimal Transport problem and jointly learns neuron-to-expert assignments with routing policies via Straight-Through Estimators.The framework achieves a superior sparsity–performance trade-off compared to heuristic clustering and structured pruning methods.
  • Future Work: Future work will explore data-driven affinity-matrix initialization using weight correlations or pre-computed activation statistics to accelerate Sinkhorn convergence and yield tighter clusters.The current affinity matrix A is initialized randomly.
  • Future Work: Future work will investigate hard pruning of experts with consistently low training utilization to reduce the model’s memory footprint.The proposed approach would permanently remove these experts.

Impact Statement … A.1. Single-Layer Reconstruction

DOT-MoE is evaluated as an output-aware alternative to dense-to-MoE assignment methods through controlled single-layer reconstruction analysis. Across LLaMA-2 and LLaMA-3, it achieves the lowest reconstruction error under matched expert and routing configurations while supporting more efficient deployment.

  • Impact Statement: The paper targets improved LLM inference efficiency by reducing deployment computation, energy consumption, and foundation-model carbon footprint.
  • A. The Assignment-Routing Gap: The method’s central claim is that existing dense-to-MoE approaches optimize inadequate FFN-output proxies, whereas DOT-MoE is output-aware.
  • A. The Assignment-Routing Gap: The evaluation separates assignment quality from training-recipe effects using single-layer reconstruction and full-pipeline ablation analyses.
  • A.1. Single-Layer Reconstruction: Single-layer analysis measures mean squared error between the original dense FFN output and each sparse MoE approximation.
  • A.1. Single-Layer Reconstruction: Layer 31 of LLaMA-2-7B and LLaMA-3-8B uses D = 128 neurons per expert and top-k = 10 routing, with WikiText calibration and WikiText-2 evaluation.
  • A.1. Single-Layer Reconstruction: Over 35× higher reconstruction error occurs for LLaMA-MoE v1 than DOT-MoE on LLaMA-2, compared with nearly 9× for v2 and more than 2× for CMoE.
  • A.1. Single-Layer Reconstruction: DOT-MoE consistently achieves the lowest MSE across LLaMA-2 and LLaMA-3, while random and proxy-based assignments produce substantially higher error.
  • A.1. Single-Layer Reconstruction: The results indicate that input-side statistics or intermediate activations alone are insufficient for clustering neurons to preserve FFN outputs.

A.2. Output-Aware Assignment Ablation … B.2. Numerical Stability

The ablation isolates OT-based assignment from training-recipe effects, while implementation details specify consistent hyperparameters and numerical-stability measures for Sinkhorn assignment. These measures include log-space computation, FP32 assignment logits, and temperature annealing.

  • A.2. Output-Aware Assignment Ablation: Ablation experiments use identical fine-tuning conditions for CMoE and DOT-MoE on Qwen2.5-7B with E=8 and top-k=2.The comparison uses the same data and training steps and evaluates hidden representations against the dense teacher.
  • A.2. Output-Aware Assignment Ablation: DOT-MoE preserves dense residual-stream and logit geometry substantially better than CMoE under the identical fine-tuning pipeline.Table 5 reports cosine similarity and MSE before and after the LM head.
  • A.2. Output-Aware Assignment Ablation: The combined reconstruction and training-dynamics analyses attribute the improvements to OT-based assignment rather than a stronger training recipe.This conclusion combines the single-layer reconstruction analysis with Figure 2.
  • B.1. Hyperparameters: Hyperparameters for MoEfication and supervised fine-tuning are summarized separately and kept consistent across model families unless otherwise noted.These settings are provided in Tables 6 and 7.
  • B.2. Numerical Stability: Sinkhorn iterations run in log-space to avoid underflow when τ is small and remain stable with thousands of neurons and experts.Algorithm 1 computes the soft assignment matrix Msoft through log-domain updates.
  • B.2. Numerical Stability: Assignment logits A remain in FP32 regardless of model dtype, while router weights Wr use the model’s native dtype.This precision choice is intended to stabilize Sinkhorn iterations.
  • B.2. Numerical Stability: The Sinkhorn temperature is linearly annealed from τstart = 1.0 to τend = 0.1 during warmup, shifting from exploration to sharper assignments.Validation always uses τend.

C. Effect of Training Sparsity … G. Extension to Attention Layers

The paper analyzes sparsity training, defines the alignment objective and masked sparse computation, then shows specialized, balanced experts and extends DOT-MoE to attention heads. Attention-only evaluation improves over random assignment by +17.9 points on average.

  • C. Effect of Training Sparsity: Models trained at 75% FFN sparsity generalize better to extreme inference sparsity, while 50% and 75% training converge at 30% inference sparsity.The ablation trains Qwen2.5-7B models at 50% and 75% FFN sparsity and evaluates four inference sparsity levels.
  • D. Training Objective: The training objective jointly optimizes neuron affinities and router weights using KL divergence, cross-entropy, router Z-loss, and load-balancing loss.These components respectively preserve dense-model outputs, language modeling, router numerical stability, and uniform expert utilization.
  • D. Training Objective: The router Z-loss limits excessively large logits, while the load-balancing term discourages expert collapse by encouraging uniform utilization.The load-balancing formulation uses token fractions fe and average routing probabilities p̄e.
  • E. Sparse MoE Computation: Sparse computation masks dense intermediate activations through RM⊤, combining token-to-expert routing with neuron-to-expert assignments.With k active experts and s neurons per expert, each token uses only k · s of dffn neurons while training retains the original dense weights.
  • E. Sparse MoE Computation: After alignment, the binary assignment partitions dense FFN weights into E expert modules compatible with existing sparse inference frameworks.The partition slices Wgate and Wup rows and Wdown columns according to each expert’s neuron cluster.
  • F. Expert Specialization and Utilization: t-SNE visualizations show well-separated activation clusters, indicating that converted experts specialize in processing distinct input types.The visualization is reported for layer 9 of Qwen2.5-7B, with each color representing an expert.
  • F. Expert Specialization and Utilization: Expert routing remains well balanced across most layers on WikiText-2 at 50% sparsity, with no evidence of severe expert collapse.The study measures token allocation across all transformer layers for Qwen2.5-7B.
  • G. Extension to Attention Layers: +17.9 points on average (64.1 vs. 46.2) is achieved by OT-based attention head assignment over random assignment at 50% attention sparsity.Attention-only DOT-MoE uses Nh=28, Eattn=14, sh=2, and kattn=7; DOT-MoE also outperforms LTE by +6.0 points on average at matched FFN sparsity.

H. Training Overhead · I. Additional Dense-to-MoE Baselines · I.1. Comparison with LTE

DOT-MoE adds limited training overhead, with Sinkhorn contributing little to the total cost, while its constant-compute top-k routing outperforms LTE at matched FFN sparsity. The comparison also highlights DOT-MoE’s predictable per-token computation for fused-MoE serving.

  • H. Training Overhead: DOT-MoE’s alignment phase is profiled on 8×H100 against a standard dense forward/backward pass.The profile targets the cost of Sinkhorn iterations and straight-through estimation.
  • H. Training Overhead: ∼2% of total forward-and-backward time comes from Sinkhorn iterations.This quantifies Sinkhorn’s contribution during the alignment phase.
  • H. Training Overhead: ∼15% overhead per training step is added by all DOT-MoE-specific operations over standard dense training.Most of this overhead comes from hard-assignment matrix construction rather than Sinkhorn iterations.
  • I. Additional Dense-to-MoE Baselines: The current CPU-based greedy rounding incurs additional CPU-to-GPU transfer overhead.The passage identifies this implementation choice as part of the current overhead profile.
  • I.1. Comparison with LTE: DOT-MoE is evaluated against LTE at matching FFN sparsity on LLaMA-2-7B with E=86 experts.LTE controls sparsity with scalar η and uses sigmoid thresholding to activate a variable number of experts per token.
  • I.1. Comparison with LTE: +6.0 points on average is achieved by DOT-MoE over LTE at lower FFN sparsity, 25% versus 29%.The comparison is reported at the stated sparsity levels.
  • I.1. Comparison with LTE: DOT-MoE’s softmax top-k routing keeps compute per token constant, unlike LTE’s variable expert activation from sigmoid routing.Constant per-token compute is compatible with standard fused-MoE serving kernels, whereas LTE’s routing produces unpredictable per-token compute.

I.2. Positioning Among Dense-to-MoE Methods

DOT-MoE is positioned against parameter-preserving, activation-agnostic dense-to-MoE methods with fixed active parameters per token and softmax top-k routing. CMoE and LLaMA-MoE(-v2) are the primary baselines because they share this problem setting, while other methods differ structurally or architecturally.

  • Structural positioning: DOT-MoE’s primary baselines match its parameter-preserving, activation-agnostic, fixed-active-parameters-per-token setting with softmax top-k routing.These properties are compatible with standard MoE serving frameworks.
  • Controlled comparison: DOT-MoE consistently outperforms CMoE across all three architectures in a controlled same-granularity comparison.The comparison uses E=8, top-k=2, and 1.2B fine-tuning tokens at CMoE’s own granularity.
  • Structural positioning: Read-ME is excluded as a primary baseline because it inflates parameters by 2.4×, converting a 7B dense model into a 17B MoE.This places Read-ME in the upcycling rather than parameter-preserving category.
  • Structural positioning: MoEfication and DejaVu are excluded because they target ReLU-based encoder architectures and do not transfer to SwiGLU decoder LLMs.Their ReLU-specific activation patterns address a different architectural setting.

J. Controlled Same-Granularity Comparison · K. Scalability · K.1. Scaling to 32B Parameters

DOT-MoE remains superior to CMoE under matched expert granularity and active-parameter budgets, while scaling to Qwen2.5-32B with a substantial benchmark-average advantage. Its alignment phase trains fewer than 2% of model parameters, leaving the frozen dense forward pass as the dominant cost.

  • J. Controlled Same-Granularity Comparison: DOT-MoE is evaluated against CMoE at matched granularity, using E=8, top-k=2, and 1.2B fine-tuning tokens across three model families.The controlled comparison is reported in Table 12.
  • J. Controlled Same-Granularity Comparison: +3.6, +14.3, and +2.3 are DOT-MoE’s improvements over CMoE on Qwen2.5-7B, LLaMA-2-7B, and LLaMA-3-8B, respectively.These gains are reported at CMoE’s own granularity.
  • J. Controlled Same-Granularity Comparison: 67.7 vs. 67.2 on Qwen2.5-7B shows that DOT-MoE at E=8 nearly matches DOT-MoE at E=148 on the same active-parameter budget.This indicates the advantage is not an artifact of finer expert granularity.
  • J. Controlled Same-Granularity Comparison: 64.1 avg versus 67.2 shows that CMoE at DOT-MoE’s default granularity, E=148, still trails DOT-MoE on Qwen2.5-7B.The controlled-comparison conclusion therefore holds in both granularity directions.
  • K. Scalability: Under 2% of model parameters are trained during DOT-MoE’s alignment phase, which updates only assignment logits and router weights.The frozen dense model’s forward pass dominates cost at scale.
  • K.1. Scaling to 32B Parameters: 25% active parameters is the evaluated setting for DOT-MoE on Qwen2.5-32B, with zero-shot performance measured on common-sense reasoning benchmarks.This scalability experiment is summarized in Table 13.
  • K.1. Scaling to 32B Parameters: +34.3 points is DOT-MoE’s benchmark-average improvement over CMoE at 32B parameters, with scores of 73.1 vs. 38.8.The result supports the scalability of the OT-based assignment as model scale increases.

K.2. Robustness to Sequence Length

DOT-MoE’s per-token neuron assignment and routing are independent of sequence length, enabling direct application to longer contexts. On WikiText-2, it maintains an approximately 2 PPL advantage over CMoE across context lengths up to 32K tokens.

  • K.2. Robustness to Sequence Length: Per-token neuron-to-expert assignment and routing operate independently of sequence length, allowing DOT-MoE to apply directly to longer contexts.Long-document evaluation splits documents exceeding the context window into rolling windows with sliding overlap.
  • K.2. Robustness to Sequence Length: DOT-MoE’s advantage over CMoE remains consistent across maximum sequence lengths from 2K to 32K.This robustness is evaluated on the document-level WikiText1 split using rolling log-likelihood.
  • K.2. Robustness to Sequence Length: DOT-MoE maintains a consistent ∼2 PPL improvement over CMoE across all context lengths up to 32K tokens.The evaluation uses WikiText-2 word perplexity on Qwen2.5-7B with 1.2B fine-tuning tokens and varying maximum context windows.
Loading 2606.01666v1…