Toward Calibrated Mixture-of-Experts Under Distribution Shift

Forum topic · 小凯 · 2026-06-21

Summary

This paper studies how calibration in Mixture-of-Experts (MoE) models behaves under distribution shift, focusing on how routing mechanisms interact with expert-level calibration. The authors, Gina Wong, Drew Prinster, and Suchi Saria, prove that in hard-routing MoE models, enforcing calibration at the individual expert level is sufficient to keep the overall model calibrated under a broad class of distribution shifts. However, this guarantee fails for soft-routing models. To address this gap, they propose an adversarial reweighting method that penalizes calibration error of the routing-aggregated predictions under distribution shift. Theoretically and empirically, the method improves the accuracy-calibration trade-off across model classes, prediction tasks, and both average data and hard subsets under distribution shift. Calibration aligns predicted uncertainty with empirical outcome frequencies, which is essential for interpreting and trusting reported probabilities; this work clarifies the conditions under which calibration benefits MoE models. Paper: arXiv:2506.16245.

Overview

Field: Machine Learning Authors: Gina Wong, Drew Prinster, Suchi Saria Published: 2026-06-20 arXiv: 2506.16245

Abstract (translated)

Calibration aligns a model's predictive uncertainty with the empirical frequency of its outcomes, and is critical for interpreting and trusting reported probabilities. Recent work has shown that enforcing calibration at the level of individual predictors can improve the accuracy and calibration of ensembles, with Mixture-of-Experts (MoE) models in particular demonstrating strong empirical improvements. However, the conditions under which calibration helps MoE models remain unclear. This paper studies the behavior of MoE models under distribution shift, focusing on how routing mechanisms interact with expert-level calibration.

The authors prove that in hard-routing models, expert calibration is sufficient to ensure the overall model remains calibrated under a broad class of distribution shifts. In contrast, this guarantee does not hold for soft-routing models.

To address this, the paper proposes adversarial reweighting, which penalizes the calibration error of routing-aggregated predictions under distribution shift. The authors show that this approach improves the accuracy-calibration trade-off on both the average data and hard subsets, across model classes, prediction tasks, and types of distribution shift.

Key Contributions

Theoretical analysis of how routing (hard vs. soft) affects the propagation of expert-level calibration to the full MoE model under distribution shift.
A provable calibration guarantee for hard-routed MoE models under broad shift classes.
An adversarial reweighting technique to restore calibration for soft-routed MoE models.
Empirical validation of improved accuracy-calibration trade-offs across diverse settings.

---

*Auto-collected on 2026-06-21*