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

AI summary

Overview

Research area: Bayesian deep learning and uncertainty quantification for Mixture-of-Experts (MoE) transformers, sitting at the intersection of probabilistic machine learning, efficient foundation-model fine-tuning, and model reliability.

Technical level: Advanced. The paper assumes familiarity with variational inference, the ELBO, KL divergence, graphical models, and the internals of sparse MoE routing, though its central idea (put uncertainty into the router, not the weights) is described in accessible terms.

Scope: The paper introduces Variational Mixture-of-Experts Routing (VMoER), a family of lightweight Bayesian fine-tuning methods that confine inference to an MoE layer's expert-selection stage, and evaluates them on three open MoE backbones for calibration, out-of-distribution detection, routing stability, and compute overhead.

What This Paper Is About

Modern foundation models reach trillions of parameters by sparsely activating subsets of "expert" subnetworks, with a small deterministic router deciding which experts each token visits. That router is brittle: it is sensitive to input perturbations and numerical precision, and the models built on it produce overconfident predictions whose errors are hard to detect. The paper asks whether principled uncertainty can be added at foundation-model scale without the prohibitive cost of conventional Bayesian methods, and answers by treating the routing decision itself as the latent variable to be inferred rather than the model's weights.

Key Contributions

1. A latent-variable formalisation of MoE routing. The authors recast the standard deterministic router (logits → Softmax → Top-K) as a probabilistic generative model with latent logits l, probabilities p, and a discrete selection mask z. This framing lets them reinterpret existing stabilisation heuristics — load-balancing regularisation and auxiliary losses — as implicit Bayesian priors.

2. Two inference strategies for the router. They introduce Logit-Space Inference via the Variational Gaussian Logit Router (VGLR), in both mean-field (VGLR-MF) and full-covariance (VGLR-FC) variants, and Selection-Space Inference via the Variational Temperature Scaling Router (VTSR), which learns an input-dependent temperature to reshape the decision boundary.

3. Empirical reliability gains across architectures. VMoER is reported to raise routing stability under noise by 38%, reduce in-distribution calibration error (ECE) by up to 94%, and improve out-of-distribution AUROC by 12% across Granite-MoE, Qwen-MoE, and DeepSeek-MoE.

4. Negligible compute cost. The authors verify that VMoER incurs less than 1% additional FLOPs and a small activation-memory footprint relative to weight-space Bayesian baselines.

Main Findings

• Deterministic routing is miscalibrated and brittle. The MAP Top-K baseline shows high ECE on OpenBookQA — 0.252 on Granite-MoE, 0.127 on Qwen-MoE, and 0.168 on DeepSeek-MoE — confirming overconfidence in standard routers. Motivation experiments also show routing decisions are sensitive to small random input noise.

• Low-temperature sampling sits in a "sweet spot." The motivational experiment replacing hard Top-K with sampling at temperature T < 1 is reported to improve both calibration and accuracy relative to the baseline.

• Full-covariance logit inference gives the best calibration. VGLR-FC reaches ECE of 0.015 on Granite-MoE (down from 0.252), 0.014 on Qwen-MoE (down from 0.127), and 0.054 on DeepSeek-MoE (down from 0.168), while keeping accuracy close to the MAP baseline (0.740, 0.802, and 0.800 respectively). The authors attribute this to modelling expert correlations through off-diagonal covariance terms that mean-field methods miss.

• Heuristic temperature scaling trades accuracy for calibration. The fixed-temperature Temp-Scale baseline improves accuracy on Qwen-MoE (0.804 to 0.817) but degrades it on Granite-MoE (0.746 to 0.716), which the authors trace to replacing deterministic argmax with stochastic sampling and to relying on a single global temperature.

• Internal router variance is a better OoD signal than gate entropy. On Granite-MoE, VGLR-FC's Inf-Logit-Var signal improves average AUROC from 0.659 (its own Gate-Ent) to 0.749, and MCDR's MC-Logit-Var improves from 0.655 to 0.717. The largest single gains appear on far-domain shifts: 0.844 on MedMCQA and 0.834 on MMLU-Law for VGLR-FC.

• For VTSR the interaction matters, not the temperature alone. The raw temperature signal (Inf-Temp) performs poorly at 0.509 average AUROC, while VTSR's Gate-Ent reaches 0.743 — suggesting uncertainty arises from the interaction between learned ambiguity and structural expert conflict.

• Stochastic routing stabilises expert selection under perturbation. On Qwen-MoE at noise σ = 0.010, MAP Jaccard similarity falls to 0.532 while VGLR-FC reaches 0.612 and VTSR 0.614.

• Overhead is small. On Granite-3B-MoE (L=10, S=35, H=D/4), VGLR-FC adds 1.15% activation memory and 1.07% FLOPs (0.0096 GFLOPs), VGLR-MF adds 0.78% and 0.77%, and VTSR adds 0.74% and 0.67%. Weight-space baselines add 2.61% activation memory and 2.32% FLOPs. Section 5.5 describes VGLR-FC's activation-memory cost as roughly 1.2%.

• Not reported in the main text. The tables shown cover OpenBookQA; the paper states that full results for the other in-distribution datasets (ARC-Challenge, SciQ, MedMCQA) are in Appendix D.2, full AUPRC values in Appendix D.3, and full stability results in Appendix D.4. Those appendices are not included in the provided content.

Methodology in Plain English

The authors leave the experts and the rest of the transformer untouched and intervene only where tokens are assigned to experts.

First, they reframe what a router does as a chain of uncertain events: an input token produces scores for each expert, those scores become a probability distribution, and that distribution produces a discrete choice of which experts to activate. Standard routers collapse this chain to a single deterministic path, discarding any notion of confidence.

For Logit-Space Inference, a small neural network runs alongside the original router and predicts a distribution over the scores rather than a single set of scores. It is trained as a residual on top of the frozen deterministic scores, so the model starts from the pretrained behaviour and learns a correction. The prior is centred on the deterministic solution, which simplifies the KL term to a distance between the learned residual and zero. At inference the model draws several samples and averages their Softmax outputs before applying Top-K. Because the number of experts is small (N ≤ 64), the authors can afford a full covariance matrix, letting them capture correlations between experts that a mean-field approximation would treat as independent.

For Selection-Space Inference, they avoid multi-sample averaging altogether. Instead, a lightweight network predicts a single scalar temperature per input, which rescales the original fixed logits. Low temperature sharpens the distribution toward one-hot selection; high temperature flattens it toward uniform. A regularisation term of −log T pushes the model toward higher stochasticity unless the data likelihood demands a sharp decision, and as temperature approaches zero the stochastic Sample-K operation converges back to deterministic Top-K. Gradients through the discrete sampling step are handled with Gumbel-Softmax.

They then fine-tune three pretrained MoE backbones — Granite (3B), Qwen (2.7B), and DeepSeek (16B) — and compare against a deterministic MAP router, a fixed global temperature heuristic, and two weight-space Bayesian baselines (MC Dropout and SWAG) using accuracy, NLL, ECE, MCE, AUROC, AUPRC, and Jaccard similarity of expert selection under input noise.

Why This Matters

Impact on research. The paper argues that weight-space Bayesian inference is fundamentally indirect for routing: because uncertainty in router weights must propagate through a linear projection before it reaches the decision variables, it produces a posteriors whose profile is hard to control. Targeting the decision manifold instead is more direct, cheaper, and empirically better calibrated. It also supplies a Bayesian reading of long-standing engineering tricks — load balancing and entropy penalties become priors rather than ad hoc fixes — and it positions internal routing variance as an epistemic signal available before token generation, in contrast to post-hoc, output-centric methods such as verbalised confidence or semantic entropy.

Real-world applications:

• High-stakes question answering in medicine, law, and science, where a model must signal when to abstain or defer; the paper's OoD evaluation uses exactly these domains (MedMCQA, MMLU-Law, ARC).
• Open-world deployment under distribution shift, where inputs differ from training data and a reliable uncertainty score governs whether a prediction is trusted.
• Serving infrastructure with latency and memory budgets, since VTSR adds no sample-count-dependent FLOPs and the other variants stay near 1% overhead.
• Low-precision or noisy inference and fine-tuning pipelines, where routing drift and selection instability have been documented as failure modes.

Industry relevance. The paper notes that every non-proprietary foundation model in the top fifty of LMArena's text leaderboard uses an MoE architecture. Since VMoER is a fine-tuning-time intervention on a small router module that leaves expert weights and the rest of the model unchanged, it is compatible with existing MoE serving stacks and does not require retraining the backbone.

Future Directions

• Scaling beyond the tested sizes. The largest backbone evaluated is DeepSeek-MoE at 16B parameters; whether the calibration and OoD gains hold at the trillion-parameter scale the paper motivates is left open.
• Closing the accuracy gap. VGLR-FC keeps accuracy near MAP but does not exceed it on Granite-MoE or DeepSeek-MoE, and Temp-Scale shows accuracy can be lost when deterministic argmax is replaced. Preserving or improving accuracy while gaining calibration remains an open problem.
• Better uncertainty signals for VTSR. Its temperature signal alone performed poorly (0.509 average AUROC) while its entropy signal was strong (0.743), so how to expose temperature-derived uncertainty directly is unresolved.
• Broadening the evaluation. The main tables shown cover OpenBookQA and OoD detection on Granite-MoE; full results across the other datasets, architectures, and AUPRC are deferred to appendices, and the truncated conclusion leaves the authors' own summary of limitations unstated in the provided content.

Target Audience

Researchers and engineers working on Bayesian deep learning, uncertainty quantification for large language models, and MoE architecture design will get the most from this paper. It is also relevant to practitioners responsible for deploying foundation models in high-stakes or distribution-shifted settings who need calibration and abstention signals at acceptable compute cost, and to those studying fine-tuning stability in sparse routing. Readers without a background in variational inference will find the high-level framing accessible but the derivations demanding.

Authors’ 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.