How Sparse Probability Maps Shape Mixture-of-Experts Routing

Tomás Brogueira; Marcos Treviso; Miguel Couceiro

cs.LGOct 5, 2026

How Sparse Probability Maps Shape Mixture-of-Experts Routing

Authors: Tomás Brogueira, Marcos Treviso, Miguel Couceiro

Organizations: Técnico, Universidade de Lisboa · INESC-ID · Instituto de Telecomunicações · ELLIS Unit Lisbon · Gandara AI

Abstract

Mixture-of-experts (MoE) routers typically apply softmax to the router scores and keep the top-K experts, making every token use exactly K experts. Sparsity-inducing probability maps such as sparsemax, alpha-entmax and normmax can adaptively assign exact zeros to selected experts, and therefore appear to offer token-dependent expert participation, even when using the same top-K machinery. In this work, we study whether and how this sparsity survives training. We train matched 300M and 1B top-2 MoE language models with softmax, 1.5-entmax, sparsemax and 2-normmax, and find that the maps behave very differently once trained: at 1B, entmax discards 30% less probability mass than softmax while almost never dropping a selected expert, sparsemax retains the most mass, and normmax routes 21% of tokens to a single expert. These outcomes are not properties of the maps alone. Each map drops a selected expert only when the gap between the two largest scores reaches a fixed threshold, and the trained routers differ in the score distribution they learn: the entmax router learns scores with roughly half the spread of softmax's, which keeps its top-2 gaps below its threshold, while sparsemax and normmax, which share the same threshold, learn different gap distributions and hence different participation. Routers thus co-adapt their scores to the map, and a map's capacity to produce zeros does not by itself determine expert participation. While none of the sparse maps improves validation loss over softmax, they make the trained models far less sensitive to selecting more experts at inference: sparsemax trained with K=2 loses 0.02 nats when run with K=8, where softmax loses 0.58. Our results indicate 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.

Figures & tables

Panel of Figure 1 — Figure 1: Left: in a top- $2$ MoE router, the probability map $\pi$ is the only component we change. Middle: on the same router scores, the maps differ in how much probability mass falls outside the selected pair (discarded mass $\delta$ , gray) and in whether a selected expert receives zero weight (a filler , dashed), which lowers the number of contributing experts $K_{+}$ ; shown for two illustrative $8$ -expert score vectors whose top-2 gaps are $0.5$ and $1.1$ . Right: statistics of our trained 1B MoE models.

Model	Val.	LMB (ppl)	LMB	ARC-E	ARC-C	CSQA	HS	OBQA	PIQA	SQA	WG	Avg.
300M params.
Dense	3.411	44.8	31.8	46.7	28.6	37.2	30.0	31.8	60.3	42.3	50.0	39.9
[.4pt/2pt] Softmax	3.534	104.2	24.3	42.7	27.6	36.5	29.4	29.6	61.9	41.4	50.8	38.2
$1.5$ -entmax	3.536	102.8	24.2	43.0	27.1	34.6	29.4	27.8	60.0	41.7	51.7	37.7
Sparsemax	3.538	104.0	24.6	42.7	26.5	36.0	27.8	28.4	59.8	40.3	52.5	37.6
Normmax- $2$	3.536	105.9	23.4	42.8	25.5	33.3	29.6	25.2	60.5	40.9	50.5	36.9

Table 1: Predictive quality of the 300M and 1B models: DCLM validation loss, perplexity and accuracies (%) on LAMBADA (LMB) and the eight OLMES tasks (ARC-Easy, ARC-Challenge, CommonsenseQA, HellaSwag, OpenBookQA, PIQA, SocialIQA, WinoGrande). The dense references have a larger active-parameter budget.

Figure 2: Discarded probability mass (a) and number of contributing experts per token (b) during training of the 1B MoE models, measured on training batches. The horizontal axis is the number of training tokens (log scale). All maps sparsify early in training; entmax then returns to two contributors while sparsemax and normmax do not. The 300M curves are provided in Appendix C.4 .

Map	$\bar{\delta}$	$\overline{K}_{+}$	$s{=}1$ %	$s{=}2$ %	$\mathbb{E}[\delta\mid s{\geq}3]$	Nominal CV $\downarrow$	Contrib. CV $\downarrow$	Contrib. SD $\downarrow$
300M params.
Softmax	0.566	2.000	0.000	0.000	0.566	0.404	0.404	0.016
$1.5$ -entmax	0.311	1.997	0.295	0.388	0.314	0.400	0.400	0.015
Sparsemax	0.099	1.904	9.598	28.174	0.160	0.434	0.450	0.017
Normmax- $2$	0.146	1.769	23.144	16.228	0.241	0.367	0.404	0.014
[.4pt/2pt] 1B params.

Table 2: Routing statistics of the trained models. Besides discarded mass

\bar{\delta}

and activated experts

\overline{K}_{+}

, we also report the fraction of token–layer states with singleton support and with support of exactly two experts, the discarded mass conditional on support larger than two, the coefficient of variation of the expert loads counting every selection or only positive-weight selections (Contrib. CV), and the standard deviation of the contributor loads (Contrib. SD). Lower CV and SD indicate more even expert loads. Bold marks the lowest value at each scale, including ties at the displayed precision. Rows under softmax-based load balancing means that we used softmax probabilities for the auxiliary balancing loss.

Figure 3: Top-2 score gaps of the trained 1B routers. Each panel pools

285{,}107

targets from each of the four MoE layers with equal weight. Gray dashed lines mark the singleton thresholds

g=1

for sparsemax and normmax-2 and

g=2

for

1.5

-entmax. Softmax has no finite singleton threshold. Panel annotations give the fractions at or above each threshold, computed on the full sample, including density beyond

g=4

for softmax. The peaks near

g=1

for sparsemax and normmax-2 are consistent with gradient saturation: for

g>1

, these maps are locally constant and gradients through the routing probabilities vanish.

Map	$K{=}1$	$K{=}2$	$K{=}4$	$K{=}8$	$\Delta$	$K{=}1$	$K{=}2$	$K{=}4$	$K{=}8$	$\Delta$
	300M					1B
Softmax	3.703	3.560	3.623	3.913	$+0.354$	3.360	3.254	3.328	3.835	$+0.582$
$1.5$ -entmax	3.711	3.562	3.598	3.693	$+0.131$	3.363	3.261	3.306	3.512	$+0.250$
Sparsemax	3.715	3.565	3.566	3.568	$\mathbf{+0.004}$	3.365	3.263	3.274	3.284	$\mathbf{+0.021}$
Normmax- $2$	3.703	3.562	3.574	3.586	$+0.024$	3.368	3.266	3.286	3.329	$+0.064$

Table 3: Robustness to the inference-time budget. NLL (nats/target) on the retained DCLM diagnostic sample when a model trained with

K{=}2

is run, with frozen weights, with

K\in\{1,2,4,8\}

selected experts;

\Delta

is the increase from

K{=}2

K{=}8

. Retention and contributor counts for every setting are in Appendix C.3 .

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Layers	Width	Q/KV	width	layers	Experts	Total	Active	Router
			Heads	FF	MoE		Parameters
Dense $283$ M	10	768	12/3	3072	–	–	283M	283M	–
MoE $300$ M	8	320	8/2	2048	4	26	297M	108M	33K
Dense $1$ B	18	1536	24/6	6144	–	–	1.01B	1.01B	–
MoE $1$ B	9	768	12/3	3072	4	27	1.01B	302M	83K

Appendix

Table 4: Architectures and parameter counts. FF width is the hidden width of each SwiGLU branch; MoE layers replace every second feed-forward block. Nominal active parameters include the full embedding tensors and count two experts per MoE layer, even when one combine weight is zero.

	300M			1B
Map	$g\geq 1$ (%)	$g\geq 2$ (%)	Own (%)	$g\geq 1$ (%)	$g\geq 2$ (%)	Own (%)
Softmax	45.702	13.785	0.000	44.757	7.824	0.000
$1.5$ -entmax	15.347	0.295	0.295	8.258	0.048	0.048
Sparsemax	9.598	0.001	9.598	6.883	0.006	6.883
Normmax- $2$	23.144	0.133	23.144	21.256	0.014	21.256

Appendix

Table 5: Fraction of top-2 score gaps at least

1

(the sparsemax and normmax-2 threshold) and at least

2

(the

1.5

-entmax threshold) for every trained router, averaged over layers, next to the router’s own singleton rate. For each sparse router the own rate equals the fraction at its own threshold; softmax has no finite threshold.

	300M, MoE layer				1B, MoE layer
Map	2	4	6	8	2	4	6	8
Singleton rate $\Pr(s{=}1)$ (%)
Softmax	0.000	0.000	0.000	0.000	0.000	0.000	0.000	0.000
$1.5$ -entmax	0.059	0.048	0.072	1.000	0.003	0.001	0.031	0.159
Sparsemax	27.934	5.451	2.002	3.003	19.267	0.887	3.573	3.806
Normmax- $2$	58.281	18.621	9.147	6.526	54.455	11.980	11.483	7.104

Appendix

Table 6: Routing statistics of every MoE layer (network depths 2, 4, 6 and 8) on the Paloma WikiText-103 test documents. Each block reports one statistic of Table 2 per layer instead of its layer average.

		300M					1B
Map	Threshold	RMS	$g_{50}$	$g_{90}$	$g_{99}$	Singletons (%)	RMS	$g_{50}$	$g_{90}$	$g_{99}$	Singletons (%)
Softmax	–	0.98	0.91	2.14	3.15	0.00	0.86	0.90	1.88	2.92	0.00
$1.5$ -entmax	$2$	0.52	0.45	1.12	1.71	0.30	0.36	0.44	0.95	1.45	0.05
Sparsemax	$1$	0.41	0.36	0.84	1.27	9.60	0.32	0.40	0.81	1.20	6.88
Normmax- $2$	$1$	0.53	0.64	1.12	1.60	23.14	0.38	0.61	1.06	1.42	21.26

Appendix

Table 7: Learned score spread and top-2 gap distribution of the trained routers on the Paloma WikiText-103 test documents. RMS is the root mean square of the centered scores and

g_{50}

g_{90}

g_{99}

are percentiles of the top-2 score gap

z_{(1)}-z_{(2)}

, all averaged over the four MoE layers. A selected expert is dropped exactly when the gap reaches the map’s threshold, so the singleton rate is the fraction of gaps at or above it.

	300M, MoE layer				1B, MoE layer
Map	2	4	6	8	2	4	6	8
RMS
Softmax	0.875	1.082	1.037	0.935	0.818	0.847	0.861	0.902
$1.5$ -entmax	0.467	0.531	0.552	0.518	0.355	0.354	0.379	0.352
Sparsemax	0.384	0.398	0.455	0.416	0.322	0.283	0.336	0.339
Normmax- $2$	0.500	0.539	0.554	0.532	0.406	0.348	0.374	0.373

Appendix

Table 8: Score scale and top-2 gap percentiles of every trained router in every MoE layer (network depths 2, 4, 6 and 8) on the Paloma WikiText-103 test documents; Table 7 averages these values over layers. RMS is the root mean square of the centered scores and

g_{q}

the

q

th percentile of the top-2 gap.

	300M, budget $K$				1B, budget $K$
Map	1	2	4	8	1	2	4	8
Loss (nats)
Softmax	3.703	3.560	3.623	3.913	3.360	3.253	3.328	3.835
$1.5$ -entmax	3.711	3.562	3.598	3.693	3.363	3.261	3.306	3.512
Sparsemax	3.715	3.565	3.566	3.568	3.365	3.263	3.273	3.284
Normmax- $2$	3.703	3.562	3.574	3.586	3.368	3.266	3.286	3.329

Appendix

Table 9: Frozen-weight inference at budgets

K\in\{1,2,4,8\}

on the retained DCLM diagnostic sample (

489

blocks,

1{,}000{,}983

scored targets) for models trained at

K=2

: loss (nats per target), retained mass

1-\bar{\delta}

and number of contributors

\overline{K}_{+}

Figure 4: Discarded probability mass (a) and number of contributing experts (b) during training of the 300M MoE models, measured on training batches; the horizontal axis is the number of training tokens (log scale).

Figure 5: Top-2 score gaps of the original 300M seed-42 routers on the same 61 Paloma WikiText-103 test documents. Samples, normalization and threshold lines follow Figure 3 . The peaks near

g=1

for sparsemax and normmax-2 are likewise consistent with gradient saturation beyond the singleton threshold.

Figure 6: Top-2 score gaps of the 300M sparsemax and

1.5

-entmax routers trained with the softmax-based balancing loss. Each panel pools

285{,}107

target positions from each of the four MoE layers with equal weight, using the same 61 Paloma WikiText-103 test documents as Figure 5 . Dashed lines mark the singleton threshold

g=1

for sparsemax and

g=2

for

1.5

-entmax, with

g=1

also shown in the entmax panel for reference.

Explore similar work

Jun 1, 2026cs.LG

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Mixture-of-Experts (MoE) models scale by activating only a small subset of experts per token. However, training such models remains challenging because top-

k

routing is discrete and non-differentiable, requiring gradient estimators for expert selection whose design remains a central open problem. We introduce ProbMoE, a probabilistic routing framework that models expert selection as a distribution over cardinality-constrained expert subsets and formulates routing as probabilistic inference in this discrete subset space. We first propose ProbMoE Exact-

k

routing, which samples

k

-expert subsets in the forward pass, and the backward pass uses gradients through each expert's exact marginal probability as a tractable surrogate for the true gradient. ProbMoE naturally generalizes to a dynamic-

k

routing setting, where both training and inference constrain the routing cardinality to the same predefined range, allowing adaptive expert allocation per token. Across benchmarks and model backbones, ProbMoE Exact-

k

achieves strong performance compared to competitive baselines, with improved expert utilization and routing diversity; ProbMoE Dynamic-

k

achieves comparable performance with fewer activated experts.

Heng Zhao, Zilei Shao, Guy Van den Broeck +1

Department of Computer Science, University of Virginia, Charlottesville, USA · Department of Computer Science, University of California, Los Angeles, Los Angeles, USA

Large Language Model Routing Probabilistic Inference Cardinality

Ask PDF

May 8, 2026cs.LG

When Are Experts Misrouted? Counterfactual Routing Analysis in Mixture-of-Experts Language Models

Mixture-of-Experts (MoE) language models route each token to a small subset of experts, but whether the routes selected by a trained top-

k

router are good ones is rarely evaluated directly. Holding the model fixed, we compare each standard route against sampled equal-compute alternatives for the same token and score each by the next-token probability it assigns to the realized token in a verified reasoning trajectory. The result is sharply token-conditional: the standard router is well-aligned with route utility on confident tokens but uninformative on the fragile tokens that drive hard reasoning, where lower-loss equal-compute routes consistently exist inside the frozen model but are not selected. The same pattern holds across Qwen3-30B-A3B, GPT-OSS-20B, DeepSeek-V2-Lite, and OLMoE-1B-7B, and follows structurally from how standard top-

k

training evaluates routing decisions: the language modeling loss scores only the executed route, and load balancing depends only on aggregate routing statistics. A minimal router-only update to the final-layer router, leaving every expert and every other router frozen, is sufficient to shift pass@K on AIME 2024+2025 and HMMT 2025 for both Qwen3-30B-A3B and GPT-OSS-20B, suggesting that at least part of the failure reflects router-reachable misallocation rather than expert capacity alone.

Youngsik Yoon, Siwei Wang, Wei Chen +1

Department of Computer Science and Engineering, POSTECH, South Korea · Microsoft Research Asia, Beijing, China · Graduate School of Artificial Intelligence, POSTECH, South Korea

Mixture-Of-Experts Experts Large Language Model Routing Deepseek Loss Function

Ask PDF

Apr 1, 2026cs.AI

Self-Routing: Parameter-Free Expert Routing from Hidden States

Mixture-of-Experts (MoE) layers increase model capacity by activating only a small subset of experts per token, and typically rely on a learned router to map hidden states to expert assignments. In this work, we ask whether a dedicated learned router is strictly necessary for MoE routing. We propose Self-Routing, a parameter-free routing mechanism that uses a designated subspace of the token hidden state directly as expert logits, eliminating the router projection entirely while leaving the rest of the MoE layer unchanged. We evaluate Self-Routing on language modeling across different expert counts and model scales, and on ImageNet-1K classification by comparing it against a standard learned router, random-routing baselines, and dense non-MoE baselines. Our results show that Self-Routing remains competitive with the learned-router baseline while removing all dedicated routing parameters, and yields more balanced expert utilization, with about 17 % higher average normalized routing entropy and no explicit load-balancing loss. On ImageNet-1K with DeiT-S/16, Self-Routing also slightly improves over the corresponding learned-router MoE. These findings suggest that effective MoE routing can emerge from the hidden representation itself without requiring a separate learned router module.

Jama Hussein Mohamud, Drew Wagner, Mirco Ravanelli

Mila - Quebec AI Institute · Universite de Montreal · Concordia University

Mixture-Of-Experts Large Language Model Routing Hidden States

Ask PDF