How Sparse Probability Maps Shape Mixture-of-Experts Routing

What it claims

Whether a sparse probability map actually reduces how many experts a token uses depends on the scores the router learns, not just the map. Each map has a fixed singleton threshold on the gap between the two largest router scores, g = z(1) − z(2). At or above it, the whole mass goes to one expert and the second selected expert gets zero weight (a "filler").

• 1.5-entmax: threshold g ≥ 2.
• Sparsemax and normmax-2: threshold g ≥ 1, so on identical scores they always agree on one vs two contributors.
• Softmax: no finite threshold.

All three sparse routers learn much tighter score spreads than softmax: RMS 0.32–0.38 against 0.86. Entmax's gaps therefore stay below its threshold of 2. Sparsemax and normmax share a threshold but learn different gap distributions, so they end up with 6.9% vs 21.3% single-expert tokens.

How it works

The paper separates two outcomes using the support size s (the number of experts with positive probability) under top-2 selection:

• E[δ] = Pr(s≥3) · E[δ | s≥3]. Mass is discarded only when more than two experts are positive.
• E[K+] = 2 − Pr(s=1). Participation drops only on a singleton.

So "retains more mass" and "uses fewer experts" are different regimes.

• Entmax sits almost entirely in the s≥3 regime: 99.9% of states have s≥3.
• Normmax sits in the s=1 regime.
• Sparsemax has the most exact-two supports (14.1%), which keep all mass with two contributors.

Training curves show all sparse maps drop contributors early. Entmax then recovers to about 2.000 contributors, while sparsemax settles at 1.93 and normmax at 1.79.

Setup:

• Backbone: Llama 3-style, with every second FFN replaced by an MoE layer.
• 300M model: 8 layers, width 320, 4 MoE layers × 26 experts.
• 1B model: 9 layers, width 768, 4 MoE layers × 27 experts, 302M active parameters.
• Routing: K=2 with Mixtral-style selection, no capacity limit and no token dropping.
• Load balancing: Switch-style loss with coefficient 0.01, computed on the full map probabilities. Fillers count in the selection frequencies, and an expert with p=0 gets no balancing gradient from that token.
• Training: DCLM-baseline, 10B tokens (19,073 steps × 524,288 tokens), AdamW, cosine schedule, bf16.
• Routing diagnostics: 61 WikiText-103 documents from Paloma at context 2,048, giving 285,107 targets per layer.

Evidence

1B models:

The dense reference, with 3.3× the active parameters, scores 3.103 / 45.0. At 300M the validation losses are 3.534–3.538. Per-task leaders vary across maps, with no consistent winner.

• Inference-time K: running the K=2-trained 1B models at K=8 costs sparsemax 0.02 nats, against 0.58 for softmax.
• Balancing ablation (300M): computing the balancing loss on softmax probabilities instead of the sparse map's probabilities makes balance worse. Sparsemax's contributor CV rises from 0.450 to 0.668, and entmax's from 0.400 to 0.485.

How to use it

No code yet ("will be released upon acceptance"). Practical takeaways from the results:

• Expect sparsemax to cost load balance; it has the highest contributor CV at both scales.
• Compute the balancing loss on the sparse map's own probabilities, not softmax probabilities.
• To check whether a trained router will actually drop experts, measure its top-2 gap distribution against the map's threshold.

Limits

• No quality gain: no sparse map beats softmax on validation loss or average accuracy.
• Narrow configuration: only top-2 routing, α=1.5, γ=2, a single balancing coefficient, two small scales and 10B tokens.
• No compute numbers: the fetched text reports no wall-clock or FLOP savings from fillers.
• Truncated page: the text stops mid-§4.3. The K=8 result is available only as the abstract/intro numbers, without the §4.4 detail.
• Preprint: this is a v1 preprint.