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

• URL: http://arxiv.org/abs/2603.09453v2
• Date: Sun, 15 Mar 2026 15:47:04 GMT
• Title: Variational Routing: A Scalable Bayesian Framework for Calibrated Mixture-of-Experts Transformers
• Abstract summary: We introduce Variational Mixture-of-Experts Routing (VMoER), a structured Bayesian approach for modelling uncertainty in MoE layers.<n>VMoER confines inference to the expert-selection stage which is typically done by a deterministic routing network.<n>Across tested foundation models, VMoER improves routing under noise by 38%, and increases out-of-distribution AUROC by 12%.
• Score: 5.534626267734823
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Architecture-agnostic Lipschitz-constant Bayesian header and its application to resolve semantically proximal classification errors with vision transformers [0.0]
  This work presents an architecture-agnostic Lipschitz-constant Bayesian header that can be integrated into feature extractors such as vision transformers.<n>We also propose a novel metric to capture uncertainty and confidence across misclassification rates, as well as an adaptive arithmetic-mean fusion scheme.<n>Although computational costs increase due to Monte Carlo sampling, the method offers plug-and-play compatibility with pre-trained backbones.
  arXiv  Detail & Related papers  (2026-05-07T09:18:06Z)
• Too Correct to Learn: Reinforcement Learning on Saturated Reasoning Data [55.84428098924793]
  We propose Constrained Uniform Top-K Sampling (CUTS), a parameter-free decoding strategy enforcing structure-preserving exploration.<n>We integrate this into Mixed-CUTS, a training framework synergizing exploitative and exploratory rollouts to amplify intra-group advantage variance.<n> Notably, Mixed-CUTS improves Pass@1 accuracy on the challenging AIME25 benchmark by up to 15.1% over standard GRPO.
  arXiv  Detail & Related papers  (2026-04-20T16:43:28Z)
• Sharp Convergence Rates for Masked Diffusion Models [53.117058231393834]
  We develop a total-variation based analysis for the Euler method that overcomes limitations.<n>Our results relax assumptions on score estimation, improve parameter dependencies, and establish convergence guarantees.<n>Overall, our analysis introduces a direct TV-based error decomposition along the CTMC trajectory and a decoupling-based path-wise analysis for FHS.
  arXiv  Detail & Related papers  (2026-02-26T00:47:51Z)
• Rethinking Test-Time Training: Tilting The Latent Distribution For Few-Shot Source-Free Adaptation [3.5808917363708743]
  We study test-time adaptation of foundation models for few-shot classification under a completely frozen-model regime.<n>We propose arguably the first training-free inference method that adapts predictions to the new task by performing a change of measure over the latent embedding distribution induced by the encoder.
  arXiv  Detail & Related papers  (2026-02-02T18:17:29Z)
• Adapformer: Adaptive Channel Management for Multivariate Time Series Forecasting [49.40321003932633]
  Adapformer is an advanced Transformer-based framework that merges the benefits of CI and CD methodologies through effective channel management.<n>Adapformer achieves superior performance over existing models, enhancing both predictive accuracy and computational efficiency.
  arXiv  Detail & Related papers  (2025-11-18T16:24:05Z)
• Bayesian Mixture of Experts For Large Language Models [2.889541910837398]
  We present a post-hoc uncertainty estimation framework for large language models (LLMs) based on Mixture-of-Experts architectures.<n>Bayesian-MoE applies a structured Laplace approximation to the second linear layer of each expert, enabling calibrated uncertainty estimation.<n> Experiments on common-sense reasoning benchmarks with Qwen1.5-MoE and DeepSeek-MoE demonstrate that Bayesian-MoE improves both expected calibration error (ECE) and negative log-likelihood (NLL) over baselines.
  arXiv  Detail & Related papers  (2025-11-12T04:24:20Z)
• Bayesian Mixture-of-Experts: Towards Making LLMs Know What They Don't Know [0.0]
  Mixture-of-Experts (MoE) has enabled the creation of massive yet efficient Large Language Models.<n>Standard deterministic routing mechanism presents a significant limitation.<n>This thesis proposes a structured textbfBayesian MoE routing framework.
  arXiv  Detail & Related papers  (2025-09-28T12:07:35Z)
• Diffusion Models for Solving Inverse Problems via Posterior Sampling with Piecewise Guidance [52.705112811734566]
  A novel diffusion-based framework is introduced for solving inverse problems using a piecewise guidance scheme.<n>The proposed method is problem-agnostic and readily adaptable to a variety of inverse problems.<n>The framework achieves a reduction in inference time of (25%) for inpainting with both random and center masks, and (23%) and (24%) for (4times) and (8times) super-resolution tasks.
  arXiv  Detail & Related papers  (2025-07-22T19:35:14Z)
• Federated Bayesian Deep Learning: The Application of Statistical Aggregation Methods to Bayesian Models [0.9940108090221528]
  Aggregation strategies have been developed to pool or fuse the weights and biases of distributed deterministic models. We show that simple application of the aggregation methods associated with FL schemes for deterministic models is either impossible or results in sub-optimal performance.
  arXiv  Detail & Related papers  (2024-03-22T15:02:24Z)
• Ensemble Kalman Filtering Meets Gaussian Process SSM for Non-Mean-Field and Online Inference [47.460898983429374]
  We introduce an ensemble Kalman filter (EnKF) into the non-mean-field (NMF) variational inference framework to approximate the posterior distribution of the latent states. This novel marriage between EnKF and GPSSM not only eliminates the need for extensive parameterization in learning variational distributions, but also enables an interpretable, closed-form approximation of the evidence lower bound (ELBO) We demonstrate that the resulting EnKF-aided online algorithm embodies a principled objective function by ensuring data-fitting accuracy while incorporating model regularizations to mitigate overfitting.
  arXiv  Detail & Related papers  (2023-12-10T15:22:30Z)
• Distributionally Robust Models with Parametric Likelihood Ratios [123.05074253513935]
  Three simple ideas allow us to train models with DRO using a broader class of parametric likelihood ratios. We find that models trained with the resulting parametric adversaries are consistently more robust to subpopulation shifts when compared to other DRO approaches.
  arXiv  Detail & Related papers  (2022-04-13T12:43:12Z)
• Deblurring via Stochastic Refinement [85.42730934561101]
  We present an alternative framework for blind deblurring based on conditional diffusion models. Our method is competitive in terms of distortion metrics such as PSNR.
  arXiv  Detail & Related papers  (2021-12-05T04:36:09Z)
• Uncertainty Estimation Using a Single Deep Deterministic Neural Network [66.26231423824089]
  We propose a method for training a deterministic deep model that can find and reject out of distribution data points at test time with a single forward pass. We scale training in these with a novel loss function and centroid updating scheme and match the accuracy of softmax models.
  arXiv  Detail & Related papers  (2020-03-04T12:27:36Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.