How do sparse probability maps shape mixture-of-experts routing after training?

Why exact zeros do not survive as expected

A mixture-of-experts router scores the experts, a probability map turns those scores into weights, and a top-K operator keeps the K highest-weighted experts, whose outputs are combined with renormalized weights. With softmax, every token uses exactly K experts and the unselected mass is simply discarded by truncation. Sparsity-inducing maps such as sparsemax, alpha-entmax and normmax can assign exact zeros, so a selected expert may receive zero weight and drop out of the mixture, making the number of contributing experts data-dependent even when the budget K stays fixed.

The paper asks whether that capacity actually produces token-dependent participation once the router is trained. The complication is that router scores are learned: changing the map also changes the scores the map receives, so a trained sparsemax router and a trained softmax router differ both in how they transform scores and in which scores they learned to produce. A map that can output zeros may therefore rarely do so after training.

Three measures that are easy to conflate

The study separates three quantities that are often treated as one. Retention is the probability mass kept by the selected experts; the discarded mass of a token is the mass of the unselected experts, and because renormalization moves exactly that mass onto the selected experts, it also equals the total variation distance between the routing distribution and its selected mixture. Participation is the number of selected experts with non-zero weight, written K+, which is at most K and equals 1 under top-2 selection whenever a selected expert has zero weight. Predictive quality is validation loss and downstream accuracy.

The paper also distinguishes two load histograms over experts: the nominal load counts every selection, fillers included, while the contributor load counts only positive-weight selections, each assignment adding 1/K. Because removing fillers also changes the mean, the two conventions can rank routers differently, so a balance claim only means something once the convention is named.

Matched models isolate the probability map

The authors train matched 300M and 1B top-2 MoE language models with softmax, 1.5-entmax, sparsemax and 2-normmax under an identical architecture, data, seed and balancing objective, so the probability map is the only component that changes. Training uses DCLM-baseline with the Meta-Llama-3 tokenizer, a nominal 10B-token budget over 19,073 updates, and a Switch-style balancing loss with coefficient 0.01 averaged over the four MoE layers; there is no router z-loss. The four maps add no parameters at fixed alpha and gamma.

Routing diagnostics record the router scores of every trained model and measure their spread and top-2 gaps relative to each map's sparsification threshold. The checkpoint diagnostic runs every model in FP32 on 61 Paloma test documents, giving 285,107 targets per layer and 1,140,428 target-layer events per model. Targets and layers are dependent observations from one checkpoint per map, not training replications.

Each map drops an expert at a fixed gap

The mechanism that makes the comparison tractable is that each map has a fixed sparsification threshold on the scores alone. A selected expert is dropped exactly when the top-2 score gap reaches 2 for 1.5-entmax and reaches 1 for sparsemax and normmax-2. The singleton rate of a trained model is therefore nothing other than the fraction of its learned gaps that reach its map's threshold, which is also a check of the recorded scores against the recorded routing.

Because the threshold is a condition on the scores, a router can change how often it drops an expert either by reshaping its standardized gaps or by moving its score scale relative to the threshold. Multiplying the scores by a positive factor multiplies every gap and leaves the shape of the gap distribution unchanged, so the singleton rate becomes the probability that the scaled gap reaches the threshold. Nothing in ordinary backpropagation pins that scale, which is why the paper treats it as a learned quantity.

Trained routers learn different score scales

At 1B, all three sparse routers learned markedly smaller score spreads than the softmax router, with a centered score RMS between 0.32 and 0.38 against 0.86 for softmax. For 1.5-entmax the consequence is decisive: its 99th-percentile gap is 1.45, well below its threshold of 2, so the trained router produces a singleton in only 0.05% of states, even though 7.8% of the softmax router's gaps reach that threshold. The 1.5-entmax router therefore discards 30% less probability mass than softmax while almost never dropping a selected expert.

Sparsemax and normmax-2 share the threshold 1, so on any given score vector they always agree on whether a token keeps one or two contributors, yet the trained models differ by 0.14 contributors per token. That difference can only come from the scores the two routers learned: the median gap of the normmax-2 router is 0.61 against 0.40 for sparsemax, and its gaps reach 1 three times as often. Normmax-2 routes 21% of tokens to a single expert, while sparsemax retains the most mass of the three sparse maps.

Participation is depth-dependent, not uniform

Layer averages hide a strong depth dependence. Normmax-2 routes 54% of tokens to a single expert in the first MoE layer but only 7% in the last; sparsemax shows the same pattern with smaller magnitude, from 19% to 4%; 1.5-entmax has essentially no singletons anywhere. Early layers learn the sparsest routers, which suggests that the useful degree of participation is itself depth-dependent and gives a further reason not to set K arbitrarily.

The paper also tests whether the balancing objective explains the differences. It trains 300M 1.5-entmax and sparsemax routers with softmax used only in the auxiliary loss, while routing and expert combination still use the sparse maps. In both runs the load CV rises, meaning traffic becomes less evenly distributed across experts, and the singleton fraction falls. Per-token participation and load balance therefore move in opposite directions, which does not support the simple explanation that sparsemax's imbalance comes solely from zero-probability experts receiving no direct balancing gradient.

Sparse maps trade quality for inference robustness

None of the sparse maps improves validation loss over softmax, and the paper notes that the small quality differences between maps, at most 0.012 nats of validation loss, do not establish a ranking given a single training seed per configuration. What the sparse maps do change is sensitivity to the inference-time budget. When models trained with K=2 are run with frozen weights at K=8, softmax loses 0.58 nats at 1B because it is forced to spread weight over six experts it never learned to combine, while sparsemax loses only 0.02 nats, since it keeps assigning zero weight to most of the extra slots and uses 4.3 experts on average at K=8. Normmax-2 and 1.5-entmax fall in between, at 0.064 and 0.250 nats.

The robustness is one-sided. Reducing the budget to K=1 costs every map about 0.10 nats, and for every map the loss is lowest at the trained budget, so selecting more experts to retain more mass does not help either. Retention, participation and quality are three distinct quantities, and the paper's own framing is that adaptive MoE routing has to be designed around the joint behavior of the probability map and the learned scores rather than around the map alone.

What to control in a sparse router

The practical priority that follows from these results is the score scale, not the choice of map. Because the router co-adapts its scale to the map, that scale is currently left to the optimizer, and the paper argues it would be better learned explicitly, as is already done for sparse attention. A router whose scale is under control could then be trusted to choose its own budget: the sparse maps already assign zero weight to most of the extra slots when K is raised, so selecting every expert in the support instead of a fixed top-K, capped by an at-most-k operator, is a small step from what the trained models already do.

The comparison to keep in view is between sparsemax and normmax-2, which share the same threshold of 1 and therefore make identical decisions on any single score vector, yet end up 0.14 contributors per token apart purely because their routers learned different gap distributions. That gap, not the map's nominal capacity to output zeros, is the quantity a designer can actually influence. The paper's limitations qualify how far this reaches: contributor counts measure expert participation, not compute, so whether zero-weight experts are skipped and what that saves depends on the dispatch implementation, which the study does not measure, and the score-scale analysis is descriptive rather than an identification of the training mechanism that produced the smaller scales.

Sources used in this answer