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

Toward Calibrated Mixture-of-Experts Under Distribution Shift

arXiv: 2606.20544v1

论文信息

标题: 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

3 分钟速览

  • 研究问题：研究混合专家模型在数据分布发生偏移时的概率校准问题。论文发现，即使每个专家模型本身已经校准，软路由机制下的整体预测仍可能因为专家组合方式的微妙变化而严重失准。
  • 核心方法：提出基于信息熵平衡的对抗性训练方法，通过加重高损失训练样本的权重，主动暴露并修补那些最容易因路由配置变化而失准的脆弱样本。
  • 关键结果：在多个基准测试的困难子集上，该方法大幅降低了校准误差。例如，在 CIFAR-10H 的困难子集上，校准误差 ECE 从 0.281 降至 0.074（表 1）；在 CivilComments 数据集上，ECE 从 0.108 降至 0.037（表 3）。
  • 主要局限：实验仅使用包含四个专家的小型模型进行验证。论文作者指出，尽管失准机制源自软路由本身，但方法在更大规模、更稀疏的模型上的效果仍有待验证。
  • 适合读者：从事混合专家模型、模型可靠性与校准、或分布外泛化研究的算法工程师和研究人员。对于需要理解模型预测概率是否可信的实践者，本文提供了重要的理论视角和实用工具。

论文背景和研究动机

现代大规模机器学习系统广泛采用混合专家（Mixture-of-Experts, MoE）架构。该架构通过一个路由器将不同输入动态分配给不同专家模型处理，从而实现条件计算——每个输入仅激活部分参数，模型容量可以增长，但计算量不成比例增加。由于专家在训练中接触不同数据子集，它们各自形成专长，使整体系统高效且强大。

一个常见的直觉是，MoE 对数据分布变化应具有天然鲁棒性。举例而言，一个拥有可靠数学专家和可靠历史专家的 MoE 系统，即使训练集偏重历史问题而测试集偏重数学问题，人们仍预期其整体保持可靠。论文明确指出，这个直觉具有误导性。问题恰恰出在软路由机制上。

在软路由中，路由器不只选择一个专家，而是生成一组权重，将输出定义为各专家预测的加权平均。这看似更灵活，但却引入了一个隐藏的脆弱性：许多不同的路由配置和专家预测组合，可能产生完全相同的聚合置信度。在训练分布上，不同配置的错误恰好相互抵消，使模型看似校准良好；但测试时，只要这些配置的出现频率发生变化——即使每个输入、每个专家的行为完全不变——整体预测的置信度就不再与其准确率匹配。

论文将这种分布偏移称为 “路由引发的重加权”，并指出其破坏性在于它温和到难以察觉：它不需要引入新输入、不需要改变条件标签概率、也不需要任何专家改变其预测。

核心方法和技术细节

论文首先从理论上对比了硬路由和软路由下的校准保持条件。在硬路由下，输入空间被分区，每个输入仅由一个专家处理。校准仅依赖于 “被选中专家的身份” 和 “该专家给出的置信度” 这两个统计量。因此，只要测试分布在每个 “专家-置信度” 切片内保留了标签的条件分布，整体校准就能保持。即使是各区域访问频率的剧烈变化，也无关紧要。

但在软路由下，情况复杂得多。每个专家看到的是被路由权重加权的重叠数据视图。由于标量聚合预测 f(x)=∑krk(x)fk(x)f(x) = \sum_k r_k(x) f_k(x) 是一个多对一映射，许多不同配置可能产生同一个预测值 pp。校准要求是在所有预测为 pp 的样例中，标签为正的频率为 pp；但这仅要求在训练集上配置级别的偏差相互抵消，并不要求每种配置本身是校准的。测试时的配置频率变化会打破这种平衡。论文在命题 4.1 中严格证明，要在任意配置重加权下都保持校准的充要条件是，每种配置的标签期望都等于聚合预测值——这是一个远比硬路由更强的条件，几乎不可实现。

基于此，论文提出了实际解决方案。既然无法直接观测导致脆弱的配置（配置本身是高维的模型内部统计量），论文转而使用一个可观测的代理信号——每样本合页损失。高损失的样本大概率正是那些配置脆弱、对校准破坏贡献最大的样本。训练目标由此构建为一个对抗性重加权问题：

min⁡θsup⁡w∈WEPtrain[w(X,Y)Lθ(X,Y)]\min_\theta \sup_{w \in \mathcal{W}} \mathbb{E}_{P_{\text{train}}}[w(X,Y) L_\theta(X,Y)]

其中，对抗方在一个受约束的权重族内寻找最能暴露模型弱点的样本分布。论文采用信息熵平衡方法：寻找在 KL 散度意义上与训练分布最接近、但将平均损失提升到对抗性水平的权重分布。其解析解是指数倾斜分布：

qiη=exp⁡(ηLi)∑j=1nexp⁡(ηLj)q_i^\eta = \frac{\exp(\eta L_i)}{\sum_{j=1}^n \exp(\eta L_j)}

基于此，论文设计了两个训练目标。Robust MoE 对全量批次施加此重加权，优化该分布下的期望损失。Robust Filtered 则更精细：它首先通过两个标准（混合预测是否比最佳单专家更差、专家间是否存在显著分歧）筛选出路由于敏感的样本子集，仅对该子集施加对抗性重加权，同时保留在整个批次上的 ERM 项。这样既保持了对全部数据的训练信号，又把对抗性压力集中到真正脆弱的地方。

创新点和贡献

论文的主要贡献体现在四个方面。第一，它首次从校准角度严格刻画了硬路由与软路由对分布偏移的响应差异，指出硬路由下，只要专家校准，则路由区域重加权无害；而软路由下，即使所有专家都完美校准，配置重加权仍可导致整体失准。第二，它识别出问题的根源是聚合映射的信息坍塌，以及由此造成的同一置信度内配置级错误的可抵消性。第三，论文提出了两个基于信息熵平衡的对抗性训练目标，将不可见的路由配置脆弱性转化为可优化的高损失样本重加权问题，并建立其与分布鲁棒优化、多精度审计的理论联系。第四，在图像分类、领域泛化、文本毒性检测三个不同领域的数据集上，验证了方法在人工和自然分布偏移下的校准改进效果。

实验结果分析

实验覆盖了三类数据和偏移场景：CIFAR-10H（低人类一致性图像构成的困难子集，人工偏移）、PACS（不同艺术风格域间的自然偏移）、CivilComments（提及身份群体的评论，次群体偏移）。所有方法使用相同架构——共享主干网络配上 4 个专家和一个软路由 MLP。

在 CIFAR-10H 上，所有非鲁棒基线在困难子集上的 ECE 都很高：单专家 0.276，Vanilla MoE 0.281，MoCaE 0.262。Robust MoE 将其降至 0.074（表 1）。值得注意的是，MoCaE 通过对各专家进行后处理温度缩放来单独校准，这确实略微降低了整体 ECE（从 0.055 到 0.049），但对困难子集上的改进甚微（从 0.281 到 0.262），直接印证了 “仅校准专家并不足够” 的论断。

在 CivilComments 数据集上，改进更加显著。基于 DistilBERT 的文本模型同样受益：Robust MoE 将困难子集的 ECE 从 0.108 降至 0.037，且困难子集准确率也从 0.866 提升到 0.868（表 3）。按身份子群细分的结果（图 2）表明，这些改进在各子群间普遍存在，并非由某个特定群体驱动。

在 PACS 跨域泛化实验中，鲁棒方法在全部四个目标域均取得最低 ECE（表 2），但不同目标域之间存在权衡：Robust MoE 在某些域上损失了 2-3 个百分点的准确率。恰恰是 Robust Filtered 方法在多个场景下提供了更好的准确率-校准折衷：它在 CIFAR-10H 上取得最低的整体 ECE（0.013），在 CivilComments 上达到最高准确率（0.915），同时维持大幅优于非鲁棒基线的校准水平。

温度缩放后处理（TS）并未消解训练改进。在 CIFAR-10H 困难子集上，非鲁棒方法中最好的 TS 后 ECE 是 MoCaE 的 0.241，而 Robust MoE 可达 0.115；在 CivilComments 困难子集上，Robust MoE 的 TS 后 ECE 为 0.029，显著优于 Vanilla MoE 的 0.100 和 MoCaE 的 0.091（表 1，表 3）。这表明改进并非单纯来自全局置信度重标定，训练目标本身改变了模型对不同样本的置信度分配模式。

实践建议

对于在实际系统中使用软路由 MoE 的工程师，本文提供几点可操作的指导。

评估时不要满足于整体 ECE。 论文的核心发现在于，即使整体校准误差可接受、各专家也看似校准良好，路由配置的重加权仍可能在次群体上产生严重失准。因此，应尽可能定义有意义的子群（如数据来源、难度、身份属性等），并分别检查校准指标。

后处理温度缩放无法根本解决路由引发的失准。 单一温度参数只能移动置信度的全局尺度，无法修复因配置重加权导致的同一置信度内错误率差异。论文实验中，MoCaE 对专家单独温度缩放后，困难子集 ECE 改善甚微，正体现了这一局限。

对抗性重加权的实现成本适中。 Robust MoE 或 Robust Filtered 只需在训练循环中增加一步指数倾斜权重的计算（公式 exp⁡(ηLi)\exp(\eta L_i)），计算开销可控。Robust Filtered 的筛选条件（混合预测相对于最佳专家的后悔值、专家间分歧）也完全可由现有模型输出计算，无需额外标注。

注意准确率与校准的权衡。 鲁棒方法在某些场景下会损失少量整体准确率（如 PACS 部分域中准确率下降约 2-3 个百分点），但大幅改善困难子集上的校准。Robust Filtered 是一种减轻此权衡的选择，它把对抗性压力集中在路由敏感的样本上，同时保留全量均匀训练的 ERM 锚定项。

从预热开始。 论文在所有实验中使用 ERM 预热 40% 的训练周期（如 50 轮中的前 20 轮），让模型先建立稳定特征表示，再应用对抗性重加权。消融实验（附录图 E.8）表明，过早启动会放大初始化噪声，而过晚启动则留给鲁棒训练的时间不足。