变分路由:用于校准混合专家Transformer的可扩展贝叶斯框架
Variational Routing: A Scalable Bayesian Framework for Calibrated Mixture-of-Experts Transformers
AI总结:
提出变分混合专家路由(VMoER),通过将贝叶斯推断限制在专家选择阶段,实现大规模模型的不确定性校准,在微调基础模型上显著提升路由稳定性、降低校准误差并提高分布外检测AUROC,且额外计算开销极小。
摘要:
基础模型越来越多地部署在需要理解其输出不确定性的场景中,这对于确保负责任部署至关重要。虽然贝叶斯方法为不确定性量化提供了原则性方法,但其计算开销使得在基础模型规模下进行训练或推理不切实际。最先进的模型通过精心设计的稀疏性(包括混合专家(MoE)层)实现了数万亿的参数数量。在这项工作中,我们通过引入变分混合专家路由(VMoER)展示了大规模下的校准不确定性,这是一种用于建模MoE层不确定性的结构化贝叶斯方法。VMoER将贝叶斯推断限制在通常由确定性路由网络完成的专家选择阶段。我们使用两种推断策略实例化VMoER:对路由logits的摊销变分推断和推断用于随机专家选择的温度参数。在微调测试的基础模型上,VMoER在噪声下将路由稳定性提高了38%,校准误差降低了94%,分布外AUROC提高了12%,同时额外FLOPs增加不到1%。这些结果表明,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 fine-tuning 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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。