Variational Routing: A Scalable Bayesian Framework for Calibrated MoE Transformers

Conference: ICML 2026
arXiv: 2603.09453
Code: To be confirmed
Area: Model Compression / LLM Efficiency / AI Safety
Keywords: Mixture-of-Experts, Bayesian Inference, Calibration, Uncertainty Quantification, Sparse Routing

TL;DR

The variational routing framework VMoER is proposed—performing variational inference on the routing decisions of MoE layers rather than weight inference leads to efficient Bayesian uncertainty modeling. It reduces calibration error by 94% and improves routing stability by 38% while maintaining <1% additional FLOPs overhead.

Background & Motivation

Background: Foundation model scales have reached trillions of parameters, and efficient expansion is achieved through MoE sparse expert routing. However, current routing mechanisms adopt deterministic Top-K strategies, which are prone to incorrect expert selection under input perturbations.

Limitations of Prior Work: (1) Deterministic routing is sensitive to input noise, leading to brittle failures; (2) Predictions are highly overconfident, resulting in large calibration errors; (3) Existing Bayesian methods target weight uncertainty, which leads to high computational overhead and is unsuitable for trillion-parameter scales.

Key Challenge: How to inject uncertainty awareness into MoE models with minimal computational cost to ensure reliable model deployment.

Goal: Design a lightweight Bayesian framework to directly model routing decisions (rather than weights) probabilistically.

Key Insight: Reformulate MoE routing as a latent variable model and observe that—(1) deterministic routing implicitly ignores the uncertainty chain of logits → probabilities → selection; (2) Top-K operations are essentially multi-label problems.

Core Idea: Shift from weight space to decision space for variational inference—performing probabilistic modeling directly on routing logits or temperature parameters via amortized inference, thus bypassing the complexity of high-dimensional weight posteriors.

Method

Overall Architecture

VMoER contains two complementary inference paths—(1) Logit Space Inference: Apply a variational Gaussian distribution \(q_\phi(\mathbf{l}|\mathbf{u})\) to routing logits \(\mathbf{l}\), explicitly modeling correlation between experts; (2) Selection Space Inference: Learn input-dependent temperature parameters \(T_\phi(\mathbf{u})\) to dynamically adjust decision boundaries, achieving stochastic expert selection via Sample-K instead of Top-K.

Key Designs

1. Variational Gaussian Logit Routing (VGLR):

   • Function: Performs amortized variational inference on routing logits, explicitly capturing expert correlations through full-covariance modeling.
   • Mechanism: Adopts a centered prior \(p(\mathbf{l}|\mathbf{u})=\mathcal{N}(\mathbf{l}_{det}, \mathbf{I})\), where \(\mathbf{l}_{det}=\mathbf{u}\mathbf{W}_r\). The posterior mean is \(\boldsymbol{\mu}_{post}(\mathbf{u})=\mathbf{l}_{det}+\Delta\boldsymbol{\mu}_\phi(\mathbf{u})\), where the inference network learns residual corrections rather than starting from scratch. Cholesky factorization \(\boldsymbol{\Sigma}_{post}=\mathbf{LL}^\top\) parameterizes the covariance (complexity \(O(N^2)\), which is acceptable for \(N \le 64\)). MC sampling is averaged during inference.
   • Design Motivation: Weight space inference indirectly propagates parameter noise via linear projection; directly modeling routing decision variables (logits → probabilities) is more efficient. Full covariance goes beyond the mean-field assumption to capture expert correlations.

2. Variational Temperature Scaling Routing (VTSR):

   • Function: Learns input-dependent temperature parameters \(T_\phi(\mathbf{u})\) to dynamically adjust softmax sharpness, enabling efficient single-dimensional variational inference.
   • Mechanism: Constrains the variational family to a 1D manifold—moving along the trajectory \(q_\phi(\mathbf{p}|\mathbf{u})=\text{Softmax}(\mathbf{l}_{det}/T_\phi(\mathbf{u}))\) defined by deterministic logits and input-dependent temperature. Sample-K sampling is performed via Gumbel-Softmax. KL regularization simplifies to Shannon entropy.
   • Design Motivation: VGLR requires multiple samples leading to inference latency; VTSR is restricted to the scale parameter space, with a computational overhead of only \(O(D_H)\) (<0.67% FLOPs).

3. Centered Prior and Residual Learning:

   • Function: Guarantees that pre-trained routing performance is not lost during fine-tuning by constraining the posterior near the deterministic solution.
   • Mechanism: Instead of learning the posterior from scratch, the residual \(\Delta\boldsymbol{\mu}_\phi(\mathbf{u})\) is learned and added to the original logits, making the KL term automatically regularize around zero.
   • Design Motivation: Routing often gets stuck during fine-tuning; the centered prior provides stability.

Loss & Training

VGLR: \(\mathcal{L}_{ELBO}=\mathbb{E}_{q_\phi(\mathbf{l}|\mathbf{u})}[\log p(\mathbf{y}|\mathbf{l},\mathbf{u})]-\beta D_{KL}(q_\phi(\mathbf{l}|\mathbf{u})\|\mathcal{N}(\mathbf{0},\mathbf{I}))\). VTSR: Primarily optimizes reconstruction, implicitly pushing temperature towards the prior via a proxy loss \(\mathcal{L}_{reg}=-\log T_\phi(\mathbf{u})\).

Key Experimental Results

Main Results

Dataset	Model	Metric	MAP Baseline	VGLR-MF	VGLR-FC	VTSR
OpenBookQA	Granite-3B	ECE ↓	0.252	0.026	0.015	0.052
OpenBookQA	Qwen-2.7B	ECE ↓	0.127	0.028	0.014	0.022
OpenBookQA	DeepSeek-16B	ECE ↓	0.168	0.067	0.054	0.060

Ablation Study

Experiment Item	Granite ECE	Qwen ECE	Findings
Deterministic Top-K	0.252	0.127	Baseline is overconfident
Fixed Temperature Scaling	0.107	0.102	Unstable across models (accuracy drops by 3%)
VGLR-FC Full Covariance	0.015	0.014	Calibration error reduced by 94%
Noise Robustness (σ=0.01)	Jaccard=0.532	Jaccard>0.612	VGLR stability improved by 38%
OoD Detection AUROC	0.659 (Baseline)	0.749 (VGLR)	Internal logit variance signal is superior to gating entropy

Key Findings

• Full covariance is key—explicit correlation modeling significantly improves calibration.
• VTSR outperforms global fixed temperature in accuracy stability.
• Internal inference uncertainty provides a stronger signal for OoD detection than predictive entropy.

Highlights & Insights

• Probabilistic Generative Model Perspective: Formulates MoE routing as a latent variable model, interpreting heuristic load balancing and auxiliary losses as implicit Bayesian priors.
• From Weight Space to Decision Space: Directly inferring routing logits or temperature parameters captures necessary uncertainty while avoiding the curse of dimensionality.
• Dual-Path Flexible Design: VGLR offers optimal calibration but slightly higher inference latency; VTSR sacrifices a small amount of precision for zero additional sampling cost through single-pass inference.
• Transferable Components: Centered prior + residual learning and temperature scaling 1D manifold simplification are generalizable.

Limitations & Future Work

• VTSR training instability—temperature parameters are prone to collapse, requiring careful initialization.
• Evaluation is limited to MCQA next-token prediction, failing to cover error accumulation in long-sequence generation.
• Larger scales were not evaluated—DeepSeek-16B is the largest.
• Improvements: Stable variational objectives for VTSR; expansion to sequence-level uncertainty; mixing with weight-space Bayesian methods.

Related Work & Insights

• vs Weight Space Methods (MCDropout/SWAG): The latter models the entire parameter space (2.6% FLOPs); Ours models only routing decisions (<1%).
• vs Heuristic Stabilization: Existing methods (fixed temperature, load balancing regularization) lack probabilistic interpretation; Ours learns input-dependent uncertainty.
• vs Output Space Uncertainty (Semantic Entropy): The latter aggregates output distributions post-hoc; Ours extracts epistemic uncertainty directly from internal routing decisions.

Rating

• Novelty: ⭐⭐⭐⭐⭐ First to systematically apply variational inference to MoE routing decisions instead of weights.
• Experimental Thoroughness: ⭐⭐⭐⭐ 3 SOTA architectures + multi-dimensional evaluation; limited to MCQA tasks, max 16B.
• Writing Quality: ⭐⭐⭐⭐⭐ The theory is clear, and the derivation of the probabilistic generative process is rigorous.
• Value: ⭐⭐⭐⭐⭐ Points toward an efficient path for the reliable deployment of trillion-parameter foundation models.

Related Papers

• [NeurIPS 2025] FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed MoE Training
• [ICML 2026] ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts
• [ICML 2026] Optimal Bayesian Stopping for Efficient Inference of Consistent LLM Answers
• [ICML 2026] DOT-MoE: Converting Dense LLMs to MoE with Differentiable Optimal Transport
• [ICML 2026] Do Transformers Need Three Projections? A Systematic Study of QKV Sharing Schemes