Selective Sinkhorn Routing for Improved Sparse Mixture of Experts

Duc Anh Nguyen; Huu Binh Ta; Nhuan Le Duc; Tan Minh Nguyen; Toan Tran

REVIEW 4 major objections 5 minor 1 cited by

Selective Sinkhorn Routing for Improved Sparse Mixture of Experts

T0 review · 4 major / 5 minor · reviewed 2026-08-03 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read Applying Sinkhorn optimal-transport routing for as little as 0.1% of training steps, and using the transport values as routing weights, improves sparse mixture-of-experts models without auxiliary balancing losses.

desk verdict A plausible incremental routing tweak with consistent but modest gains; the headline numbers are driven by noise injection that the paper never cleanly isolates from the Sinkhorn component. read the letter →

arxiv 2511.08972 v2 pith:IU54JO72 submitted 2025-11-12 cs.LG

Duc Anh Nguyen , Huu Binh Ta , Nhuan Le Duc , Tan Minh Nguyen , Toan Tran This is my paper

classification cs.LG

keywords sparsemixtureofexpertsoptimaltransportSinkhornroutingloadbalancingexpertentropyregularizationlanguagemodelingimageclassification

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper tries to establish that sparse mixture-of-experts (SMoE) models can be made to use experts more evenly and train better not by adding auxiliary loss terms but by occasionally (0.1%–1% of steps) computing token-to-expert assignments as entropy-regularized optimal transport and using those transport values directly as routing weights. The authors show this minimal intervention matches or beats load-balancing and z-loss on language modeling and image classification, improves robustness to corrupted inputs, and is far cheaper than running Sinkhorn every step. If true, it means a major source of training instability and objective misalignment in SMoE — auxiliary balancing objectives — can be replaced by a lightweight, principled routing mechanism.

What carries the argument

The central object is the entropy-regularized optimal transport plan Π̂ computed by the Sinkhorn-Knopp algorithm, with a cost matrix derived from gating scores (either raw or row-wise softmax). The paper uses the values of Π̂, after top-k selection and renormalization, as routing weights — rather than using Π̂ only to choose which experts receive a token. The plan is constrained so each token distributes total mass 1 and each expert receives expected mass m/n, which enforces balanced utilization during training. A stochastic 'selective' schedule invokes Sinkhorn with probability p and standard softmax routing with probability 1−p, preserving gradients to the gating matrix through the softmax

What would settle it

A controlled scale-up test: train an SMoE above 1B parameters with SSR at p=0.001 against vanilla SMoE with the same compute budget and evaluate on standard benchmarks; if the gains seen at 216M/550M parameter scale do not appear or reverse, the claimed general benefit fails. Alternatively, an ablation that removes the KL renormalization step while keeping Sinkhorn selection would isolate whether transport values, rather than Sinkhorn balancing alone, are what improve performance.

Watch

Extended reading notes

Core claim

Routing a sparse mixture-of-experts (SMoE) model via entropy-regularized optimal transport usually means using the Sinkhorn transport plan only to select the top-k experts, and running Sinkhorn at every training step. This paper claims that the transport plan values themselves should be used as routing weights, and that Sinkhorn needs to run only for a tiny fraction of training steps (0.1%–1% per epoch), with standard softmax routing at all other times and at inference. The authors prove that taking the top-k entries of a transport-plan row and renormalizing gives the sparse distribution closest to that row in KL divergence; that batch-level balancing constraints are meaningful during traini

Load-bearing premise

The load-bearing assumption is that forcing balanced expert usage during a small fraction of training steps improves the final model's task performance; the paper gives empirical support but no proof of a generalization benefit, and effect sizes are small in several settings.

Editorial extensions

If this is right

Sinkhorn routing can replace auxiliary balancing losses, removing objective misalignment and associated hyperparameter tuning.
Using transport plan values as routing weights, not just as expert selectors, is what drives the improvement; top-k renormalized transport rows are the KL-closest sparse distributions to the full plan.
Load balancing constraints are only justified during training over batches; at inference, applying Sinkhorn or noise degrades performance relative to deterministic softmax routing.
Noise added to the cost matrix during training guarantees every expert has positive selection probability and acts as implicit data augmentation; it is beneficial in training but must be disabled at inference.
SSR adds almost no overhead compared to vanilla SMoE (fraction of a percent) and is orders of magnitude faster than running Sinkhorn at every step.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

The success of such a tiny fraction suggests the mechanism may act as an occasional regularizer that prevents expert collapse without permanently distorting the learned routing; the same pattern could be tested in other discrete assignment problems (e.g., attention heads, retrieval indices) where balance is desirable.
Because the Sinkhorn branch does not backpropagate into the gating matrix, the method implicitly relies on the softmax branch alone to learn compatibility; a natural extension is to add a smooth gradient path through the transport plan and test whether performance improves further.
The robustness gains on corrupted inputs are consistent with balanced routing acting as a regularizer that yields more evenly specialized experts; measuring representation collapse or per-expert specialization during training would test this interpretation.
The paper's own result that inference-time Sinkhorn or noise hurts suggests that any load-balancing mechanism is best confined to training; this could inform future MoE designs that use auxiliary losses only during early or sparse phases.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

Quick take: SSR is a neat, lightweight training-time routing tweak — use the Sinkhorn transport plan values as routing weights, and only run Sinkhorn on 0.1–1% of steps. The vanilla SSR variants give small but consistent PPL improvements on WikiText; the larger headline numbers come from the noise variants, and the paper never cleanly isolates the noise from the Sinkhorn component. That is the main soft spot.

What's genuinely new: prior Sinkhorn routers (Kool, Clark, Liu) used the transport plan only to pick top-k experts; SSR uses the plan values as the routing weights. Sparse application is a sensible answer to the "Sinkhorn is expensive" critique, and the training-time overhead numbers make the point. The paper is also refreshingly clear that Sinkhorn and noise should be train-time only; Table 2 confirms that plain softmax at inference is best.

Where it's soft: The central claim — "minimal OT routing improves SMoE without auxiliary losses" — is not actually isolated. The w/ noise variants add Gaussian noise to the cost matrix, and the "SMoE w/ noise" baseline is not a matched control: it couples trainable noise, an auxiliary loss, and different initialization. A proper ablation would compare SSR-L/S with and without noise against vanilla SMoE with the same noise schedule. Without that, the +2.37 ImageNet gain looks like it could be mostly noise-as-regularization, not balanced assignment. Second, the theory does less work than advertised. Prop 4.2 says the training-time Sinkhorn plan is balanced over the batch; it does not show that transient balancing transfers to the final softmax-routed model, which is the actual claim. Props 4.1 and 4.3 are correct but standard. Also, after top-k truncation and renormalization in Eq. 6, the routing weights no longer satisfy the balanced column constraint C3 — so the stated motivation is not strictly present in the forward pass. Third, the vision results have no error bars and the hyperparameters (p, xi, alpha_noise) are selected per benchmark, which makes the "0.1%-1%" claim a loose summary rather than a rule.

That said, the paper is honest in its claims, the direction of the effect is consistent across many settings, and the method is simple enough to test. It deserves a serious referee. I'd recommend major revision: add a matched noise-only control, report error bars on ImageNet, and dial back the theoretical framing to what the propositions actually establish.

For what it's worth, I'd bring it to reading group as a useful example of a training-time routing trick with a clean presentation but underdetermined evidence.

Referee Report

4 major / 5 minor

Summary. This paper proposes Selective Sinkhorn Routing (SSR) for sparse Mixture-of-Experts. During training, with probability p, the token-to-expert scores are replaced by the solution of an entropy-regularized optimal transport problem with a uniform column marginal; the top-k entries of each row of this transport plan are renormalized and used directly as routing weights. With probability 1-p, standard softmax top-k routing is used. Optionally, zero-mean Gaussian noise is added to the cost matrix. At inference, deterministic softmax routing is always used. The paper reports lower perplexity on WikiText-103 and Enwik-8 and higher ImageNet accuracy/robustness than baselines, and shows very low training overhead. It also gives three propositions intended to justify the design: top-k renormalization minimizes row-wise KL, Sinkhorn enforces balanced routing during training but not at inference, and noise gives every expert nonzero selection probability.

Significance. If the empirical claim holds, SSR is attractive because it replaces auxiliary balancing losses with a tiny amount of Sinkhorn computation and adds no parameters. The experimental coverage is broad (two language datasets, ImageNet and robustness benchmarks, multi-head MoE variants) and the paper includes useful ablations of p, ξ and α_noise and reports wall-clock overhead. The theoretical component is, however, much weaker than the paper suggests: the propositions are either elementary or, in the case of Prop. 4.3, mathematically incorrect as stated. A matched ablation isolating noise from Sinkhorn balancing is missing, and the paper provides no significance testing, so the central causal claim remains underdetermined. No circular reasoning was found; the comparisons are to external baselines.

major comments (4)

[§5.2, Tables 1 and 4; Eq. (9)] The two headline results are obtained by SSR-with-noise variants: -0.551 PPL (Table 1, 1-Head SSR-S w/ noise) and +2.368 Top-1 (Table 4, SSR-L w/ noise). The baseline 'SMoE w/ noise' described in §5.1 is not a matched control, since it couples trainable gating noise with a load-balancing loss. There is no ablation in which the same Gaussian perturbation (Eq. (9), same α_noise and σ) is applied to the vanilla softmax top-k router with all other hyperparameters fixed. Without this isolation, the larger gains cannot be attributed to Sinkhorn-based assignment rather than to stochastic exploration; the central claim of the abstract ('minimal optimal transport-based routing improves SMoE performance') is therefore underdetermined by the present experiments.
[§4.2 / Appendix Prop. 5.3] The displayed product formula for P_i is derived by treating the events {g_i + α_noise ε_i > g_j + α_noise ε_j} as independent over j. They are not independent because they share the same ε_i. For example, when all g_j are equal and σ=1, the true top-1 probability is 1/n (by symmetry), whereas the product equals (1/2)^{n-1}; for n=3 this is 1/4 vs 1/3. The correct expression is an integral over ε_i. The qualitative conclusion P_i>0 remains true, so the proposition can be repaired, but as written it is false.
[§4.1 Eq. (6) and Prop. 4.2] The balanced-marginal property (C3) is proven only for the full transport plan Π̂. Once the top-k experts are selected and the weights are renormalized per token as in Eq. (6), the routing weights α actually used on Sinkhorn steps no longer have column sums m/n. Thus even during the Sinkhorn training steps the forward pass is not exactly balanced. More importantly, Prop. 4.2 only states that the full plan balances the training batch; it does not prove that applying this balancing on a small fraction p of steps transfers to a final model that routes by softmax top-k at inference. This missing link is load-bearing because the final inference model is not Sinkhorn-routed. The paper should either prove a transfer statement or explicitly reposition Props. 4.1-4.3 as motivation.
[§5.2, Tables 1 and 4] Several non-noise improvements are within the reported variance. For example, in the 4-Head clean WikiText-103 row, SSR-L improves by -0.047 PPL while the reported ± values overlap substantially (vanilla 34.284±0.056, SSR-L 34.237±0.058). No significance tests or confidence intervals are reported anywhere, and Table 4 (vision) has no error bars at all. Since the paper's conclusion is a claim of consistent improvement, the authors should provide per-seed results with significance tests or CIs, especially for the smaller effects.

minor comments (5)

[§5.2 text] 'Attached WikiText-103' should be 'Attacked WikiText-103'.
[§4.2 and Appendix] Propositions 4.1-4.3 are referred to as 'Theorem 4.2' and 'Theorem 4.3' at the end of §4.2 and in the supplement proofs; align the numbering throughout.
[Appendix §8] In the Lagrangian derivation, λ and μ are described as nonnegative with sums equal to 1. For equality-constrained problems these are free dual variables; the final scaling form is standard, but this statement is incorrect.
[Eq. (9) and figure captions] The figure captions in the supplement use 'noise = 1' to mean α_noise=1; define this once near Eq. (9) to avoid ambiguity.
[Reproducibility] No code release or seed list is mentioned; with many baselines and hyperparameter choices, a reproducibility statement would help.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity found: the derivation does not reduce to its inputs, and empirical claims rest on external baselines.

full rationale

The paper's central claim is empirical: applying Sinkhorn-based routing for a small fraction of training steps improves SMoE accuracy/perplexity compared with external baselines. No step in the derivation is equivalent by construction to the claimed outcome. The theoretical results are limited to properties of the proposed routing mechanism: Proposition 4.1 proves that the top-k renormalized transport row is the optimal k-sparse KL projection of the transport row, which is a mathematical characterization of the algorithm rather than a prediction of downstream performance. Proposition 4.2 shows only that constraint (C3) enforces uniform column marginals in expectation during training; it does not assert that balancing improves task loss. Proposition 4.3 derives positive selection probabilities under Gaussian noise, again an exploration property. None of these results assumes the performance conclusion. The paper does not invoke a load-bearing self-citation chain, uniqueness theorem, or ansatz imported from the authors' prior work; references to prior Sinkhorn routing are external and used as baselines. The empirical comparisons are against vanilla SMoE, auxiliary-loss SMoE, noisy SMoE, and Sinkhorn-based SMoE, and reported gains are benchmark outcomes rather than fitted values relabeled as predictions. Hyperparameters p, xi, and alpha_noise are tuned and reported, but tuning hyperparameters on validation does not make the headline result a fitted input called a prediction. The main legitimate concern is that the largest gains appear in the 'w/ noise' variants and the noise baseline is not a perfectly matched control, so the contribution of the Sinkhorn component is somewhat underdetermined; however, this is a correctness/ablation concern, not circularity. The reader's observed gap between the balancing theorem and the performance claim is an overreach in interpretation, but not a circular reduction.

Assumptions & free parameters 4 free parameters · 4 assumptions · 0 invented entities

The paper introduces no new physical or conceptual entities. It uses existing machinery (optimal transport, Sinkhorn algorithm, Gaussian noise) in a new configuration. The free parameters are hyperparameters, and the main unproven axiom is that balanced training yields better softmax-routed inference.

free parameters (4)

p (Sinkhorn probability) = 0.001–0.01 (best per benchmark; also tested 0.0001–0.1)
Controls the fraction of training steps where Sinkhorn routing is used. The paper tunes p on validation sets for each dataset and reports the best; the '0.1%-1%' claim depends on this choice.
ξ (entropic regularization) = 0.05–1.0 (chosen per benchmark)
Regularization strength in the optimal transport objective. It affects the sharpness of the transport plan; the choice is tuned on validation performance.
α_noise (noise scale) = 0.3–4 (chosen per benchmark)
Scales the Gaussian noise added to the cost matrix during training. This is a key hyperparameter; the largest gains in the paper come from the noise variant, and the optimal value varies across tasks.
δ (Sinkhorn tolerance) and η (iterations) = δ=0.0001, η=100
Convergence parameters for the Sinkhorn algorithm. Fixed, but they affect the precision of the transport plan and the computational cost.

assumptions (4)

domain assumption Uniform expert usage during training improves final softmax-routed model performance.
This is the core load-bearing assumption. Proposition 4.2 only shows that Sinkhorn routing enforces uniform usage over the training distribution; it does not prove that this improves task accuracy at inference, where softmax routing is used. The paper relies on this as an empirical premise.
domain assumption The gating network can be effectively trained through the softmax branch while the Sinkhorn branch provides assignments without gradients to the gating matrix.
The paper explicitly acknowledges that Sinkhorn routing decouples Wg from the computational graph (Section 4.2). It assumes that the concurrent softmax training preserves meaningful gating scores, which is necessary for the transport cost to be informative.
standard math Batch samples are i.i.d. from the data distribution, so applying the uniform marginal constraint (C3) corresponds to balancing expert usage over the population.
Used in Proposition 4.2 (training case). This is a standard law-of-large-numbers assumption.
standard math Sinkhorn-Knopp algorithm converges to the entropy-regularized optimal transport solution.
The paper relies on the standard convergence of Sinkhorn iterations (Algorithm 1, Section 3.2) to compute the transport plan.

how reviews work

0 comments

Cite this review

Pith. "Pith review of Selective Sinkhorn Routing for Improved Sparse Mixture of Experts." pith.science (2026). https://pith.science/paper/IU54JO72

@misc{pith2026251108972,
  author       = {Pith},
  title        = {Pith review of: Selective Sinkhorn Routing for Improved Sparse Mixture of Experts},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/IU54JO72}},
  note         = {Machine review of arXiv:2511.08972}
}

read the original abstract

Sparse Mixture-of-Experts (SMoE) models are scalable and computationally efficient, enabling large increases in model capacity with limited inference overhead. Existing SMoE methods often depend on auxiliary objectives, such as load-balancing loss and z-loss, or additional trainable components such as noisy gating. While these techniques encourage expert diversity, they can introduce objective misalignment, increase model complexity, or incur substantial training overhead, especially in Sinkhorn-based routing methods. In this paper, we revisit the token-to-expert assignment as an optimal transport problem. We add constraints to ensure balanced expert utilization. We show that even minimal optimal transport-based routing improves SMoE performance without requiring auxiliary balancing losses. Unlike prior approaches, our method derives gating scores directly from the transport map, leading to more balanced and effective token-to-expert assignments. Building on this insight, we introduce Selective Sinkhorn Routing (SSR), a lightweight routing mechanism that replaces complex auxiliary losses with efficient Sinkhorn-based routing while preserving flexible expert selection. Experiments on language modeling and image classification show that SSR improves training efficiency, accuracy, and robustness to input corruption.

Figures

Figures reproduced from arXiv: 2511.08972 by the authors.

**Figure 2.** Training time Overhead vs. SMoE (Vanilla) over [PITH_FULL_IMAGE:figures/full_fig_p006_2.png] view at source ↗

**Figure 3.** Performance of SSR-S on Enwik-8 under varying [PITH_FULL_IMAGE:figures/full_fig_p007_3.png] view at source ↗

Figures from the paper (3 more)

**Figure 4.** Figure 4: Training time Overhead vs. Vanilla SMoE over 2- [PITH_FULL_IMAGE:figures/full_fig_p013_4.png]

**Figure 5.** Figure 5: Training time Overhead vs. Vanilla SMoE over 4- [PITH_FULL_IMAGE:figures/full_fig_p014_5.png]

**Figure 6.** Figure 6: Training time Overhead vs. Vanilla SMoE on [PITH_FULL_IMAGE:figures/full_fig_p014_6.png]

Discussion (0). Sign in to comment.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

Mixture-of-Control: State-Aware Fine-Tuning for Transformer-based Models
cs.LG 2026-06 unverdicted novelty 5.0 of 10

Mixture-of-Control adaptively combines local and global control states in transformer fine-tuning by treating per-block states as experts in a sparse MoE setup to improve cross-block communication while keeping memory...

Reference graph

Works this paper leans on

2 extracted references · 1 linked inside Pith · cited by 1 Pith paper

[8]

Training is performed with a batch size of 48 for 80,000 steps, using a learning rate of 0.0007 with 4,000 warm-up steps and a dropout rate of 0.1

Each layer consists of a multi-head self-attention (MHA) block followed by a SMoE block, both with residual con- nections. Training is performed with a batch size of 48 for 80,000 steps, using a learning rate of 0.0007 with 4,000 warm-up steps and a dropout rate of 0.1. The model uses 8 attention heads for each MHA block and processes se- quences of 512 t...

arXiv 2024
[2024]

Selective Sinkhorn Routing for Improved Sparse Mixture of Experts

Multi-Head Mixture-of-Experts. InThe Thirty-eighth Annual Conference on Neural Information Processing Sys- tems. Zhou, Y .; Lei, T.; Liu, H.; Du, N.; Huang, Y .; Zhao, V .; Dai, A. M.; Le, Q. V .; Laudon, J.; et al. 2022. Mixture-of-experts with expert choice routing.Advances in Neural Information Processing Systems, 35: 7103–7114. Zoph, B.; Bello, I.; Ku...

arXiv 2022

Pith tools

Reviewed August 3, 2026 · model on record in the stance chip above.

[1] [8]

Training is performed with a batch size of 48 for 80,000 steps, using a learning rate of 0.0007 with 4,000 warm-up steps and a dropout rate of 0.1

Each layer consists of a multi-head self-attention (MHA) block followed by a SMoE block, both with residual con- nections. Training is performed with a batch size of 48 for 80,000 steps, using a learning rate of 0.0007 with 4,000 warm-up steps and a dropout rate of 0.1. The model uses 8 attention heads for each MHA block and processes se- quences of 512 t...

arXiv 2024

[2] [2024]

Selective Sinkhorn Routing for Improved Sparse Mixture of Experts

Multi-Head Mixture-of-Experts. InThe Thirty-eighth Annual Conference on Neural Information Processing Sys- tems. Zhou, Y .; Lei, T.; Liu, H.; Du, N.; Huang, Y .; Zhao, V .; Dai, A. M.; Le, Q. V .; Laudon, J.; et al. 2022. Mixture-of-experts with expert choice routing.Advances in Neural Information Processing Systems, 35: 7103–7114. Zoph, B.; Bello, I.; Ku...

arXiv 2022