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

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.

As AI models become more powerful, they are increasingly used to make high-stakes decisions where mistakes can be costly. However, modern "foundation models" are often overconfident, struggling to accurately communicate when they are uncertain or confused. A major reason for this is their reliance on "Mixture-of-Experts" architectures—a design that efficiently splits tasks among specialised sub-networks but uses rigid, inflexible routing systems to decide which expert handles which task. We solve this by introducing Variational Routing (VMoER), a new approach that treats this routing decision as a flexible probability problem rather than a rigid choice. By letting the AI mathematically weigh its uncertainty before choosing an expert, our method dramatically improves the model's ability to recognise unfamiliar data and avoid overconfident errors. Importantly, this solution requires almost no extra computing power. This research is a critical step toward building robust, trustworthy AI systems that not only perform well, but know the limits of their own knowledge.