Toward Calibrated Mixture-of-Experts Under Distribution Shift
2026-06-18 • Artificial Intelligence
Artificial IntelligenceMachine Learning
AI summaryⓘ
The authors study models called mixture-of-experts (MoE), which use multiple specialized predictors combined to make decisions. They focus on how well these models' predicted uncertainties match actual outcomes, especially when things change unexpectedly (distribution shift). They find that for models that pick one expert decisively (hard routing), making each expert well-calibrated leads to a well-calibrated overall model. But for models that mix experts softly, this isn't enough. To fix this, the authors propose a method that adjusts how experts are combined to improve calibration and accuracy when conditions change.
calibrationuncertaintymixture-of-expertsdistribution shiftrouting mechanismshard routingsoft routingadversarial reweightingaccuracy-calibration tradeoffensemble models
Authors
Gina Wong, Drew Prinster, Suchi Saria, Rama Chellappa, Anqi Liu
Abstract
Calibration aligns a model's predictive uncertainty with the frequencies of its empirical outcomes and is important for understanding and trusting reported probabilities. Recent work shows that enforcing calibration at the level of individual predictors can improve ensemble accuracy and calibration, with mixture-of-experts (MoE) models showing strong empirical improvements in particular; however, the conditions under which calibration helps MoE are not well understood. In this work, we study how MoE models behave under distribution shift, focusing on how routing mechanisms interact with expert-level calibration. We show that expert calibration is sufficient to ensure calibration of the overall model under a broad class of distribution shifts in hard-routed models, but is insufficient for calibrating soft-routed models. To address this, we propose an adversarial reweighting that penalizes calibration errors of the routed aggregate under distribution shift, and we demonstrate that it improves the accuracy-calibration tradeoff both on average and on difficult subsets of the data, across model classes, prediction tasks, and distribution shifts.