Variational Routing: A Scalable Bayesian Framework for Calibrated Mixture-of-Experts Transformers

Albus Yizhuo Li; Matthew Wicker

cs.LGMar 10, 2026

Variational Routing: A Scalable Bayesian Framework for Calibrated Mixture-of-Experts Transformers

Authors: Albus Yizhuo Li, Matthew Wicker

Organizations: Department of Computing, Imperial College London.

Abstract

Foundation models are increasingly being deployed in contexts where understanding the uncertainty of their outputs is critical to ensuring responsible deployment. While Bayesian methods offer a principled approach to uncertainty quantification, their computational overhead renders their use impractical for training or inference at foundation model scale. State-of-the-art models achieve parameter counts in the trillions through carefully engineered sparsity including Mixture-of-Experts (MoE) layers. In this work, we demonstrate calibrated uncertainty at scale by introducing Variational Mixture-of-Experts Routing (VMoER), a structured Bayesian approach for modelling uncertainty in MoE layers. VMoER confines Bayesian inference to the expert-selection stage which is typically done by a deterministic routing network. We instantiate VMoER using two inference strategies: amortised variational inference over routing logits and inferring a temperature parameter for stochastic expert selection. Across fine-tuning tested foundation models, VMoER improves routing stability under noise by 38%, reduces calibration error by 94%, and increases out-of-distribution AUROC by 12%, while incurring less than 1% additional FLOPs. These results suggest VMoER offers a scalable path toward robust and uncertainty-aware foundation models.

Figures & tables

Panel of Figure 1 — Figure 1: Motivation. (a) Deterministic routing is empirically brittle to input noise. (b) Simple low-temperature sampling ( $T{<}1$ ) creates a ”sweet spot” that improves both Calibration (Blue $\downarrow$ ) and Accuracy (Orange $\uparrow$ ) over the baseline (dotted). Both motivational experiments’ details can be found in Appendix B .

Figure 2: Standard deterministic router architecture.

Figure 3: MoE Routing PGM. Probabilistic graphical model formulation of MoE routing treating decisions as latent variables.

Figure 4: Weight-Space Inference. Uncertainty is modeled by sampling global parameters $\mathbf{W}_{r}\sim p(\mathbf{W}_{r}\mid\mathcal{D})$ .

Figure 5: Variational Gaussian Logit Router (VGLR) Architecture. Latent logits are sampled from a posterior parameterised by a shared-trunk network using the reparameterisation trick, then averaged to marginalise routing uncertainty.

Figure 6: Geometry of Variational Temperature Scaling. The temperature

T

acts as a radial coordinate on the simplex. A low

T

(green) pushes the distribution toward the vertices (certainty), while a high

T

(blue) pulls it toward the centre (uniformity). By learning

T_{\phi}(\mathbf{u})

, the model dynamically expands or contracts its variational search space (yellow region) based on input ambiguity.

Figure 7: Variational Temperature Scaling Router (VTSR). Learn a heteroscedastic temperature via neural network and sampling from a scaled probability for expert selection.

		Granite-MoE				Qwen-MoE				DeepSeek-MoE
Method	Type	ACC $\uparrow$	NLL $\downarrow$	ECE $\downarrow$	MCE $\downarrow$	ACC $\uparrow$	NLL $\downarrow$	ECE $\downarrow$	MCE $\downarrow$	ACC $\uparrow$	NLL $\downarrow$	ECE $\downarrow$	MCE $\downarrow$
MAP	Det.	0.746	1.38	0.252	0.472	0.804	1.37	0.127	0.273	0.802	1.120	0.168	0.321
Temp-Scale	Heur.	0.716	0.773	0.107	0.201	0.817	1.12	0.102	0.224	0.808	0.870	0.146	0.402
MCDropout	Wei.	0.734	0.650	0.037	0.298	0.798	0.982	0.064	0.152	0.794	0.729	0.103	0.302
SWAG	Wei.	0.736	0.652	0.041	0.290	0.782	0.912	0.098	0.146	0.786	0.665	0.140	0.288
VGLR-MF	Log.	0.742	0.654	0.026	0.293	0.796	0.724	0.028	0.124	0.788	0.476	0.067	0.174

Table 1: Main Results (OpenBookQA). Comparison of calibration and predictive performance across three MoE architectures on the OBQA commonsense reasoning task. Results are averaged over 5 runs. Our VGLR-FC method consistently minimises calibration error (ECE) while maintaining high accuracy. VTSR also achieves strong reliability improvements. Full results in Appendix D.2 .

		Near-Domain		Far-Domain
Method	Signal	ARC-E	ARC-C	Med	Law	Avg
MAP	Gate-Ent	0.612	0.633	0.679	0.683	0.652
MCDR	Gate-Ent	0.612	0.632	0.684	0.691	0.655
MCDR	MC-Logit-Var	0.610	0.677	0.786	0.793	0.717
VGLR-FC	Gate-Ent	0.615	0.652	0.692	0.677	0.659
VGLR-FC	Inf-Logit-Var	0.609	0.709	0.844	0.834	0.749

Table 2: OoD Detection Performance (Granite-MoE). Comparison of AUROC scores using standard gating entropy ( Gate-Ent ) versus internal router signals. Bold indicates the best signal.

Noise ( $\sigma$ )	MAP	MCDR	VGLR-FC	VTSR
0.001	0.938	0.944	0.956	0.948
0.002	0.924	0.932	0.946	0.936
0.005	0.614	0.662	0.672	0.668
0.010	0.532	0.582	0.612	0.614

Table 3: Routing Stability (Qwen-MoE). Average Jaccard Similarity of expert selection under input perturbation.

	Activation Memory			Computational Cost (FLOPs)
Method	Complexity	Params	Add. %	Complexity	GFLOPs	Add. %
Weight-Space Baselines (MCDR[parallelised], SWAGR)
All	$O(LSDN)$	20.9M	2.61%	$O(LSDN)$	0.0208	2.32%
VMoER (Ours)
VGLR-MF	$O(L(DH+HN))$	6.2M	0.78%	$O(L(DH+SN))$	0.0069	0.77%
VGLR-FC	$O(L(DH+HN^{2}))$	9.2M	1.15%	$O(L(DH+HN^{2}+SN^{2}))$	0.0096	1.07%

Table 4: Complexity and Overhead Analysis. Comparison on Granite-3B-MoE (

L{=}10

S{=}35

H{=}D/4

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Experimental setup for quantifying the brittleness of deterministic routing.

Panel of Figure 9 — Figure 9: Brittleness Analysis Results. (a) Mean Jaccard similarity across layers shows high sensitivity to minute input perturbations. (b) At a fixed diagnostic noise level ( $\gamma=0.01$ ), instability is not uniform but concentrated in specific layer groups (e.g., 0-1, 28-31).

Figure 10: Experimental setup for introducing stochastic routing at a single MoE layer.

Figure 11: Impact of stochastic routing on Accuracy (left) and ECE (right). The top row shows all layers; the bottom row excludes layer 0 to highlight trends in deeper layers. Stochasticity often improves calibration (lower ECE) without harming accuracy in later layers.

Figure 16

Specification	Granite-MoE	Qwen1.5-MoE	DeepSeek-MoE
HuggingFace ID	granite-3.1-3b-a800m	Qwen1.5-MoE-A2.7B	deepseek-moe-16b
Total Parameters	3.3B	14.3B	16.4B
Active Params	800M	2.7B	2.8B
MoE Layers	32	24	27
Hidden Dimension	1,536	2,048	2,048
Routed Experts ( $N$ )	40	60	64

Appendix

Table 5: Configurations of the base MoE models used in this study. All models are fine-tuned from their base instruct versions.

Dataset	Domain	Role 1: ID Calibration	Role 2: OoD Detection
		(Train / Val / Test)	(Shift Type)
OBQA	Commonsense Science	5000 / 50 / 500	ID Source (Training)
ARC-C	Primary Science (Hard)	2000 / 50 / 500	Near-OoD (OoD-S)
ARC-E	Primary Science (Easy)	– Not Used –	Near-OoD (OoD-S)
SciQ	Broad STEM	5000 / 50 / 500	– Not Used –
MedMCQA	Professional Medicine	5000 / 50 / 500	Far-OoD (OoD-L)

Appendix

Table 6: Benchmark datasets and their usage across experiments. Note that datasets like ARC-C and MedMCQA have explicit training splits for Exp 1 (Calibration), but are treated as held-out test sets for Exp 2 (OoD Detection) where the model is trained solely on OBQA.

		Granite-MoE				Qwen-MoE				DeepSeek-MoE
Method		ACC $\uparrow$	NLL $\downarrow$	ECE $\downarrow$	MCE $\downarrow$	ACC $\uparrow$	NLL $\downarrow$	ECE $\downarrow$	MCE $\downarrow$	ACC $\uparrow$	NLL $\downarrow$	ECE $\downarrow$	MCE $\downarrow$
Dataset: ARC-Challenge (Reasoning)
Base	MAP	0.882	0.923	0.201	0.428	0.758	1.852	0.252	0.624	0.580	1.462	0.237	0.826
Base	Temp-Scale	0.824 (.004)	0.208 (.006)	0.038 (.007)	0.284 (.003)	0.762 (.003)	1.422 (.001)	0.102 (.001)	0.304 (.005)	0.592 (.003)	1.242 (.008)	0.188 (.006)	0.646 (.003)
Weight	MCDR	0.880 (.003)	0.146 (.006)	0.028 (.003)	0.228 (.007)	0.732 (.003)	1.134 (.001)	0.198 (.011)	0.446 (.011)	0.564 (.005)	1.084 (.009)	0.192 (.023)	0.848 (.003)
Weight	SWAG	0.872 (.003)	0.138 (.006)	0.030 (.007)	0.266 (.002)	0.746 (.002)	1.122 (.006)	0.202 (.010)	0.387 (.002)	0.568 (.002)	1.183 (.001)	0.198 (.004)	0.842 (.009)

Appendix

Table 7: Main Results: Calibration & Performance by Dataset. We compare standard deterministic routing against weight-space and our proposed VMoER across three MoE architectures. Results are reported as

\text{Mean}_{(\text{std})}

over 5 runs. Our methods (VGLR, VTSR) consistently reduce calibration error without degrading accuracy.

			Near-Domain Shift				Far-Domain Shift				Average
			OBQA $\to$ ARC-E		OBQA $\to$ ARC-C		OBQA $\to$ MedMCQA		OBQA $\to$ MMLU-Law		Overall
Model	Method	Signal	AUROC	AUPRC	AUROC	AUPRC	AUROC	AUPRC	AUROC	AUPRC	AUROC	AUPRC
Granite	MAP	Gate-Ent	0.612	0.596	0.633	0.626	0.679	0.645	0.683	0.686	0.652	0.638
	MCDR	Gate-Ent	0.612	0.599	0.632	0.610	0.684	0.651	0.691	0.672	0.655	0.633
	MCDR	MC-Logit-Var	0.610	0.583	0.677	0.623	0.786	0.723	0.793	0.765	0.717	0.674
	VGLR-MF	Gate-Ent	0.622	0.603	0.642	0.622	0.682	0.637	0.673	0.664	0.655	0.632

Appendix

Table 8: OoD Detection Performance across Routing Methods and Signals. We compare the effectiveness of uncertainty signals for Out-of-Distribution detection. Bold indicates the best performance across all signals for a given model. Inf-Logit-Var consistently performs best for VGLR methods, while VTSR often benefits from Gate-Ent . Both performs equally best overall.

	Granite				DeepSeek				QWen
Noise ( $\sigma$ )	MAP	MCDR	VGLR-FC	VTSR	MAP	MCDR	VGLR-FC	VTSR	MAP	MCDR	VGLR-FC	VTSR
0.001	0.946	0.954	0.970	0.962	0.956	0.958	0.948	0.954	0.938	0.944	0.956	0.948
0.002	0.922	0.932	0.940	0.934	0.918	0.944	0.940	0.922	0.924	0.932	0.946	0.936
0.005	0.650	0.822	0.886	0.840	0.712	0.732	0.758	0.746	0.614	0.662	0.672	0.668
0.010	0.422	0.654	0.702	0.692	0.624	0.684	0.718	0.702	0.532	0.582	0.612	0.614

Appendix

Table 9: Routing Stability under Input Perturbation. We report the average Jaccard Similarity of expert selection across the modified probabilistic layers for each model.

Explore similar work

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts
Jun 1, 2026Heng Zhao, Zilei Shao, Guy Van den Broeck +1Sparse Mixture-of-Experts Expert Routing
Department of Computer Science, University of Virginia, Charlottesville, USA · Department of Computer Science, University of California, Los Angeles, Los Angeles, USA
Distributionally Robust Mixture-of-Experts Training
Oct 5, 2026Xin Teng, Muxiao Li, Hongyi WenMixture-of-Experts Language Models Sparse Mixture-of-Experts
New York University Center for Data Science, NYU Shanghai
Value-and-Structure Alignment for Routing-Consistent Quantization of Mixture-of-Experts Models
Jun 4, 2026Hancheol Park, Geonho Lee, Tairen Piao +1LLM Quantization Mixture-of-Experts Quantization
Nota Inc., South Korea.

Jun 1, 2026J/K move · Enter open · S save

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Heng Zhao, Zilei Shao, Guy Van den Broeck, Zhe Zeng

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

Sparse Mixture-of-Experts Expert Routing Mixture-of-Experts Models

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Distributionally Robust Mixture-of-Experts Training

Value-and-Structure Alignment for Routing-Consistent Quantization of Mixture-of-Experts Models