Toward Calibrated Mixture-of-Experts Under Distribution Shift

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.

AI models (and predictive models in general) should be calibrated, meaning the probabilities they predict should match the frequencies at which things actually occur. This paper studies how calibration can break down in the mixture-of-experts (MoE) setup, a popular way of building AI systems where a "router" sends each input to a set of specialized "expert" models. Unfortunately, MoE can be fragile when the test data differs from the training data. The same final prediction can be assembled from different combinations of experts, so when a shift in data changes how often each combination appears, the overall calibration can fall out of balance---and this can happen even if the expert models are individually well-calibrated. We characterize exactly when this failure occurs, and we show that it can be addressed by focusing training on difficult examples where experts disagree, as these are precisely the cases where combining experts is most unstable. We show that this focused training improves both the accuracy and calibration of MoE across a variety of image and text tasks.