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

TL;DR

Variational Mixture-of-Experts Routing (VMoER), a structured Bayesian approach for modelling uncertainty in MoE layers, is introduced by introducing VMoER, a structured Bayesian approach for modelling uncertainty in MoE layers that offers a scalable path toward robust and uncertainty-aware foundation models.

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 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 (13)

• 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 \ref{['app:motivation']}.
• 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.