Toward Calibrated Mixture-of-Experts Under Distribution Shift

Summary

This arXiv paper (2506.17580) by Gina Wong, Drew Prinster, and Suchi Saria studies calibration in Mixture-of-Experts (MoE) models under distribution shift. Calibration aligns a model's predictive uncertainty with empirical outcome frequencies, which is essential for interpreting and trusting reported probabilities. Prior work has shown that enforcing calibration at the level of individual predictors can improve the accuracy and calibration of ensemble models, with MoE showing particularly notable empirical gains—but the conditions under which calibration benefits MoE remained unclear. The authors analyze MoE behavior under distribution shift, focusing on the interaction between routing mechanisms and expert-level calibration. They prove that expert calibration is sufficient to guarantee overall model calibration for hard-routed models under a broad class of distribution shifts, but is insufficient for soft-routed models. To address this, they propose an adversarial reweighting method that penalizes the calibration error of the routing aggregate under shift, demonstrating improved accuracy–calibration trade-offs both on average and on hard data subsets, across model classes, prediction tasks, and types of distribution shift.