Bayesian Mixture-of-Experts: Towards Making LLMs Know What They Don't Know
作者: Albus Yizhuo Li
分类: cs.LG, math.ST, stat.ML
发布日期: 2025-09-28
💡 一句话要点
提出贝叶斯MoE路由框架，提升LLM不确定性感知能力
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 混合专家模型 大语言模型 贝叶斯方法 不确定性建模 模型校准
📋 核心要点
- 现有MoE模型采用确定性路由，导致模型校准不良和过度自信，无法有效识别自身知识盲区。
- 论文提出贝叶斯MoE路由框架，对路由决策建模概率分布，引入权重、logit和选择空间的不确定性。
- 实验表明，该框架显著提升路由稳定性、分布内校准和分布外检测，增强LLM的鲁棒性和自我认知。
📝 摘要（中文）
混合专家模型（MoE）架构使得构建大规模且高效的大语言模型（LLM）成为可能。然而，标准的确定性路由机制存在显著局限性：其固有的脆弱性是模型校准不良和过度自信的关键因素，导致系统常常无法识别自身知识的不足。本研究通过提出一个结构化的贝叶斯MoE路由框架来应对这一挑战。该方法不强制进行单一的确定性专家选择，而是对路由决策本身建模为一个概率分布。我们系统地研究了三种方法，在路由流程的不同阶段引入这种有原则的不确定性：在权重空间、logit空间和最终的选择空间。通过在30亿参数的MoE模型上进行的一系列受控实验，我们证明了该框架显著提高了路由稳定性、分布内校准和分布外（OoD）检测。结果表明，通过针对这一核心架构组件，我们可以创建一个更可靠的内部不确定性信号。这项工作为构建更鲁棒和具有自我意识的LLM提供了一条实用且计算上可行的途径，朝着使它们知道自己不知道什么迈出了关键一步。
🔬 方法详解
问题定义：现有MoE模型使用确定性路由机制，即每个输入token被确定性地分配给一个或几个专家。这种确定性分配方式使得模型在遇到不熟悉的输入时，仍然会强制选择一个专家进行处理，导致模型过度自信，无法准确评估自身的不确定性，从而影响模型的可靠性和泛化能力。现有方法缺乏对路由决策不确定性的建模，导致模型校准不良。
核心思路：论文的核心思路是将确定性的专家路由过程转化为概率性的路由过程。不再强制模型选择一个或几个确定的专家，而是为每个专家分配一个概率，表示该专家处理当前输入的可能性。通过对路由决策进行概率建模，使模型能够更好地表达自身的不确定性，从而提高模型的校准性能和鲁棒性。
技术框架：该论文提出的贝叶斯MoE路由框架主要包含以下几个阶段：1) 输入token经过embedding层得到输入表示；2) 路由网络根据输入表示计算每个专家的路由概率；3) 根据路由概率对专家进行加权，得到最终的输出表示；4) 使用最终的输出表示进行下游任务的预测。论文主要研究了在三个不同阶段引入不确定性的方法：权重空间、logit空间和选择空间。
关键创新：该论文的关键创新在于提出了一个结构化的贝叶斯MoE路由框架，通过对路由决策进行概率建模，引入了对专家选择的不确定性。与传统的确定性路由方法相比，该方法能够更好地表达模型自身的不确定性，从而提高模型的校准性能和鲁棒性。此外，论文系统地研究了在权重空间、logit空间和选择空间引入不确定性的不同方法，并分析了它们的优缺点。
关键设计：论文在三个空间引入不确定性的具体设计如下：1) 权重空间：对专家网络的权重进行贝叶斯建模，使用变分推断学习权重的后验分布；2) Logit空间：在路由网络的输出logit上添加噪声，例如高斯噪声或Gumbel噪声；3) 选择空间：使用Gumbel-Softmax技巧对专家选择进行软化，允许模型以一定的概率选择多个专家。论文还设计了相应的损失函数，用于训练贝叶斯MoE模型，例如KL散度损失和交叉熵损失。
🖼️ 关键图片
📊 实验亮点
在30亿参数的MoE模型上进行的实验表明，该贝叶斯MoE路由框架显著提高了路由稳定性、分布内校准和分布外检测。具体而言，该方法在分布内校准方面取得了XX%的提升，在分布外检测方面取得了YY%的提升（具体数值未知）。实验结果表明，通过对路由决策进行概率建模，可以有效提高LLM的鲁棒性和自我认知能力。
🎯 应用场景
该研究成果可应用于对可靠性要求较高的自然语言处理任务，例如医疗诊断、金融风控等。通过提高LLM的不确定性感知能力，可以降低模型在关键决策中的误判风险，提升系统的整体安全性。此外，该方法还可以用于提高LLM的泛化能力，使其在面对未知数据时能够做出更合理的预测。
📄 摘要（原文）
The Mixture-of-Experts (MoE) architecture has enabled the creation of massive yet efficient Large Language Models (LLMs). However, the standard deterministic routing mechanism presents a significant limitation: its inherent brittleness is a key contributor to model miscalibration and overconfidence, resulting in systems that often do not know what they don't know. This thesis confronts this challenge by proposing a structured \textbf{Bayesian MoE routing framework}. Instead of forcing a single, deterministic expert selection, our approach models a probability distribution over the routing decision itself. We systematically investigate three families of methods that introduce this principled uncertainty at different stages of the routing pipeline: in the \textbf{weight-space}, the \textbf{logit-space}, and the final \textbf{selection-space}. Through a series of controlled experiments on a 3-billion parameter MoE model, we demonstrate that this framework significantly improves routing stability, in-distribution calibration, and out-of-distribution (OoD) detection. The results show that by targeting this core architectural component, we can create a more reliable internal uncertainty signal. This work provides a practical and computationally tractable pathway towards building more robust and self-aware LLMs, taking a crucial step towards making them know what they don't know.