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

📄 arXiv: 2603.09453v1 📥 PDF

作者: Albus Yizhuo Li, Matthew Wicker

分类: cs.LG, cs.AI, stat.ML

发布日期: 2026-03-10

备注: 8 pages, 7 figures for main text; 16 pages for Appendix; In submission to ICML 2026;

💡 一句话要点

提出VMoER：一种可扩展的贝叶斯框架，用于校准混合专家Transformer的不确定性。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 不确定性量化 贝叶斯推理 变分推理 路由算法 基础模型 模型校准

📋 核心要点

  1. 现有基础模型缺乏有效的不确定性量化方法，限制了其在关键场景中的可靠应用。
  2. VMoER通过在MoE层的专家选择阶段引入贝叶斯推理，建模路由过程中的不确定性。
  3. 实验表明，VMoER在提升模型校准性、鲁棒性的同时，计算开销增加极小。

📝 摘要（中文）

基础模型越来越多地部署在需要理解其输出不确定性的环境中，以确保负责任的部署。虽然贝叶斯方法为不确定性量化提供了一种原则性方法，但其计算开销使得它们在基础模型规模上的训练或推理中不切实际。目前最先进的模型通过精心设计的稀疏性（包括混合专家（MoE）层）实现了数万亿的参数量。在这项工作中，我们通过引入变分混合专家路由（VMoER）来展示大规模的校准不确定性，VMoER是一种结构化的贝叶斯方法，用于建模MoE层中的不确定性。VMoER将贝叶斯推理限制在专家选择阶段，该阶段通常由确定性路由网络完成。我们使用两种推理策略实例化VMoER：对路由logits进行摊销变分推理，以及推断随机专家选择的温度参数。在经过测试的基础模型中，VMoER将噪声下的路由稳定性提高了38％，将校准误差降低了94％，并将超出分布的AUROC提高了12％，而产生的额外FLOPs不到1％。这些结果表明，VMoER为实现稳健且具有不确定性意识的基础模型提供了一条可扩展的路径。

🔬 方法详解

问题定义：论文旨在解决基础模型在实际部署中，由于缺乏有效的不确定性量化方法而导致的可靠性问题。现有方法，特别是应用于MoE模型时，计算开销巨大，难以扩展到大规模模型。确定性路由网络在专家选择时忽略了不确定性，导致模型在噪声环境下表现不佳，校准误差较高。

核心思路：论文的核心思路是将贝叶斯推理引入到MoE模型的专家选择（路由）阶段，通过对路由过程建模不确定性，提高模型的鲁棒性和校准性。VMoER避免了对整个模型参数进行贝叶斯推理，而是专注于路由logits，从而降低了计算复杂度，使其能够应用于大规模模型。

技术框架：VMoER的核心在于对MoE层中的路由网络进行改进。传统的MoE层使用确定性路由网络选择专家。VMoER则引入变分推理，将路由logits视为随机变量，并学习其分布。具体来说，论文提出了两种推理策略：一是使用摊销变分推理直接学习路由logits的分布；二是引入一个温度参数来控制专家选择的随机性，并通过推理学习该参数。整体框架保持了MoE的基本结构，只是在路由部分进行了修改。

关键创新：VMoER的关键创新在于将贝叶斯推理应用于MoE模型的路由过程，从而在不显著增加计算开销的情况下，实现了对模型不确定性的量化。与传统的确定性路由方法相比，VMoER能够更好地处理噪声数据，提高模型的鲁棒性和校准性。此外，VMoER通过限制贝叶斯推理的范围，使其能够扩展到大规模模型。

关键设计：VMoER的关键设计包括：1) 使用变分推理对路由logits进行建模，学习其分布；2) 提出两种推理策略：摊销变分推理和温度参数推理；3) 损失函数包括标准的MoE损失以及用于学习路由logits分布的KL散度项；4) 实验中，VMoER被应用于不同的基础模型，并与基线方法进行比较，以验证其有效性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，VMoER在多个基础模型上显著提升了性能。具体来说，VMoER将噪声下的路由稳定性提高了38％，将校准误差降低了94％，并将超出分布的AUROC提高了12％，而计算开销仅增加了不到1％。这些结果表明，VMoER在提高模型鲁棒性和校准性的同时，保持了良好的计算效率。

🎯 应用场景

VMoER可应用于各种需要高可靠性和不确定性量化的基础模型应用场景，例如自动驾驶、医疗诊断、金融风控等。通过提供校准的不确定性估计，VMoER可以帮助决策者更好地理解模型的预测结果，并做出更明智的决策，从而提高系统的安全性和可靠性。未来，VMoER可以进一步扩展到其他类型的模型和任务中。

📄 摘要（原文）

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.