面向分布偏移下的校准混合专家

论文信息

标题: Toward Calibrated Mixture-of-Experts Under Distribution Shift

作者: Gina Wong, Drew Prinster, Suchi Saria, et al.

发布日期: 2026-06-18

arXiv ID: 2606.20544v1

PDF 链接: 下载 PDF

研究背景与动机

机器学习模型的预测不仅需要准确，更需要对不确定性进行诚实的表达。校准（calibration）要求模型的预测置信度与实际经验频率对齐，是构建可信系统的重要基石。近年来，混合专家（Mixture-of-Experts, MoE）架构凭借条件计算带来的高效扩展能力，已成为大规模学习系统的标准组件。直觉上，若每个专家在其 “专长” 领域内均可靠，则整体预测也理应可靠；但本文作者指出，这一直觉在软路由 MoE 下存在严重误导——即便每个专家都被完美校准，只要路由配置的分布发生轻微偏移，聚合模型的校准就可能崩溃。

问题的根源在于软路由将多个专家的输出加权求和为一个标量置信度，这一运算是多对一映射：许多完全不同的路由权重与专家预测组合会映射到相同的聚合置信度。训练分布上的校准只要求这些组合的标签频率在平均意义上等于该置信度，而不强制每个组合内部都正确。一旦测试时这些组合的相对比例发生变化——无需引入新输入、改变条件标签概率或专家预测本身——校准平衡便被打破。作者将这种现象称为 “路由引起的重新加权”（routing-induced reweighting），并系统研究了硬路由与软路由在分布偏移下的校准保持条件，进而提出两种基于对抗重加权的训练目标来缓解这一问题。

硬路由的简单鲁棒性

硬路由下，每个输入被分配到一个唯一的专家，聚合预测退化为该专家的输出，输入空间被划分为不相交的路由区域。令 $h(x) = \arg\max_k r_k(x)$ 为路由选择的专家编号，统计量

S_{\text{hard}}(X) = (h(X), f_{h(X)}(X))

完全决定了模型对输入 $x$ 的预测。若每个专家在其路由区域上进行校准，即 $\mathbb{P}(Y=1 \mid S_{\text{hard}}(X)=(k,p)) = p$ 在训练分布上成立，那么只要测试分布满足

P_{\text{test}}(Y \mid S_{\text{hard}}(X)) = P_{\text{train}}(Y \mid S_{\text{hard}}(X)),

聚合模型在测试下的校准就能得到保证。重要的是，该条件允许大幅改变各路由区域出现的频率，甚至改变区域内部的协变量分布，只要不破坏专家-置信度切片内的标签分布。因此，硬路由对分布偏移具有天然的鲁棒性：校准精度被压缩到专家身份和其置信度这两个维度的瓶颈中，路由权重的重排无法穿透这一瓶颈。

软路由的脆弱性与校准失效机制

软路由则为每个输入给出一个路由分布 $(r_1(x),\dots,r_K(x))$ ，专家在重叠的、按路由权重加权的视角下进行训练。此时专家校准的条件仅约束了个体在其权重流上的边际行为：

\mathbb{E}_{P_{\text{train}}}\left[r_k(X)(Y - f_k(X)) \mid f_k(X)=p\right] = 0 \quad \forall p.

将这些条件积分并求和，仅能得到聚合残差的边际期望为零，而无法保证对每个预测水平 $p$ 的条件校准 $\mathbb{E}[Y - f(X) \mid f(X)=p] = 0$ 。究其原因，聚合预测所依赖的高维配置

S_{\text{soft}}(X) = (\{r_k(X)\}_{k=1}^K,\, \{f_k(X)\}_{k=1}^K)

被压缩为单个实数值 $f(X) = \sum_k r_k(X) f_k(X)$ 。设 $m(s) = \mathbb{E}[Y \mid S_{\text{soft}}(X)=s]$ 为配置 $s$ 的真实标签频率，则校准要求

\mathbb{E}_{P_{\text{train}}}\left[ m(S) \mid f(X)=p \right] = p.

由于 $m(s)$ 可以在同一预测水平 $p$ 上出现正、负偏差，训练分布中它们恰好相互抵消。一旦测试分布通过路由配置重加权改变了这些偏差的组合比例，校准便不再成立。

作者在 Proposition 4.1 中精确刻画了这一点：对任意只重加权 $S_{\text{soft}}$ 而保持条件分布 $Y\mid S$ 不变的测试分布，聚合模型保持校准的充要条件是，对训练分布几乎处处有 $m(S) = f(X)$ ，即每个配置内部的标签频率与它的聚合预测完全一致。该条件远比硬路由对应的条件苛刻，并且在实际的软路由 MoE 中几乎不可能自然满足。失校真正集中发生在那些多个专家同时获得高路由权重且预测明显相左的区域，在这些区域聚合置信度高度依赖路由权重的微小变化，轻微的分布偏移即可造成较大的校准残差。

对抗性重加权的鲁棒训练方法

既然直接识别并修复易碎的配置并不现实，作者转而利用可观测的严格真分数损失（proper scoring loss）作为代理变量：在路由重叠且预测冲突的例子中，聚合模型的损失通常较大。于是，训练时主动对抗性地重加权高损失样本，有望使模型在这些关键配置上变得校准稳健。

具体而言，对于每个小批次，熵平衡（entropy balancing）对手在满足期望损失提升至一定水平的前提下，选择与均匀分布 KL 散度最小的权重分布，其解为指数倾斜形式：

q_i^{\eta} = \frac{\exp(\eta L_i)}{\sum_{j=1}^n \exp(\eta L_j)}.

其中 $L_i$ 为聚合交叉熵损失， $\eta \ge 0$ 控制对抗强度。基于此，作者提出两种训练目标：

Robust MoE：将熵平衡权重直接作用于全批次损失，最小化 $\mathcal{L}_{\text{R-MoE}} = \sum_i q_i^{\eta} L_i$ 。该目标可等价写为 $\rho_{\eta}(L) + \frac{1}{\eta}\mathrm{KL}(q^{\eta}\|u)$ ，其中 $\rho_{\eta}$ 为熵风险，第二项为诱导集中惩罚，当高损失样本过度集中时增加正则化。它同时可视为 KL-ball 分布鲁棒优化的保守上界代理。
Robust Filtered：仅对路由相关子集 $A$ 应用熵平衡加权，同时保留全批次的 ERM 项。 $A$ 由两类标准筛选：(1) 聚合损失高于最佳专家损失；(2) 路由加权下的专家预测方差较大。这一设计将鲁棒压力聚焦于真正由路由引发的脆弱点，避免因非路由因素造成的高损失样本过度主导训练。

在优化过程中，权重 $q^{\eta}$ 随当前损失动态更新，梯度同时通过损失项和权重项传播，完整包含了 KL 散度项的梯度（协方差项），从而实现了对脆弱配置的平滑且稳定的压制。

实验结果与关键发现

实验覆盖图像分类（CIFAR-10H，PACS）与文本毒性检测（CivilComments），包含人工定义的困难子集和自然的分布偏移。主要对比基线包括单一专家模型、标准软路由 MoE、专家单独校准后组合的 MoCaE、以及频域感知梯度修正（FGR）等。结果揭示了三个核心发现：

鲁棒训练显著改善分布偏移下的校准：在 CIFAR-10H 的低人一致性困难子集上，Vanilla MoE 和 MoCaE 的 ECE 分别为 0.281 和 0.262，而 Robust MoE 和 FGR+Robust 组合分别降至 0.074 和 0.065；在 CivilComments 的身份提及子群上，ECE 从 0.108 降至 0.037。PACS 的四个域中，鲁棒方法几乎在全部目标域取得最低 ECE。
单独校准专家远远不够：MoCaE 只在全体数据上带来微弱的 ECE 改善（例如 CIFAR-10H 整体 ECE 从 0.055 降至 0.049），在困难子集上改善极小，与鲁棒方法之间的差距巨大。这表明混合层面的校准失衡无法通过事后缩放专家概率来修复。
Robust Filtered 在准确率-校准权衡上更优：完整对抗重加权有时会牺牲过多准确率（如 CIFAR-10H 和 PACS-Art），而聚焦路由相关样本的 Robust Filtered 在维持甚至提升准确率的同时大幅改善校准，显示出更强的实用性。

此外，可靠性图表和按置信度/人一致性分组的分析表明，非鲁棒基线即使在准确率很低的困难样本上仍给出接近 1 的置信度，而鲁棒方法将置信度分布拉宽，更贴近对角校准线。后置温度缩放虽能略微降低全局 ECE，但无法消除路由偏移带来的硬子集校准失败，进一步凸显鲁棒训练的必要性。

实践启示与未来方向

本研究对 MoE 系统的工程实践提供了几个直接指导。首先，MoE 的校准评估必须置于聚合层面，且在存在子群或分布偏移的条件下进行，仅依赖总体平均精度和专家级校准会掩盖严重失效。其次，后置全局温度缩放无法修复混合内部配置失衡，需要通过在训练阶段注入对路由敏感性的意识来解决。此外，针对软路由 MoE 的对抗重加权（尤其是 Robust Filtered）可以在少量额外计算开销下，明显提升模型在歧义或偏移数据上的可靠性，适合集成到现有的 MoE 训练流程中。

未来工作可在多个方向延伸。作者选用了紧凑的 4 专家模型以隔离聚合效应的干扰，但理论上更大规模、包含数十甚至数百个专家的稀疏 MoE 会有更丰富的路由重叠和更多的易碎配置，失校准问题可能更为严重；将本方法扩展到大规模稀疏 MoE 和生成式模型是一个自然步骤。同时，目前采用交叉熵损失作为代理，更直接利用路由权重和专家分歧构造校准正则项，或设计更具目标性的多校准约束，有望进一步提升效率。此外，探索在 MoE 路由中内置不变性表示，从结构上限制路由偏移带来的失衡，也是解决该问题的长远思路。

总结

本文系统揭示了软路由混合专家在分布偏移下校准失稳的本质原因：多对一的聚合压缩使训练分布上的偏差抵消在测试时轻易瓦解。作者从硬路由的鲁棒性入手，严格刻画了两类路由的校准保持条件，并据此设计出熵平衡对抗重加权的训练目标，迫使模型在路由重叠区域获得更一致的校准行为。实验证明了方法的有效性，且聚焦路由敏感样本的 Robust Filtered 在准确率与校准之间取得了良好平衡。整体而言，这项工作不仅深化了对 MoE 概率推断行为的理论理解，也为工程实践中部署可靠的概率预测系统提供了清晰的技术路径。