Search papers, labs, and topics across Lattice.
This paper investigates the calibration of mixture-of-experts (MoE) models under distribution shifts, revealing that while expert-level calibration ensures overall model calibration in hard-routed configurations, it fails in soft-routed setups. The authors introduce an adversarial reweighting method that penalizes calibration errors in the routed aggregate, significantly enhancing the accuracy-calibration tradeoff across various tasks and data subsets. The findings underscore the importance of routing mechanisms in achieving reliable uncertainty estimates in MoE models, particularly in challenging scenarios.
Expert calibration alone isn't enough for soft-routed MoE models; an innovative adversarial reweighting approach is needed to tackle distribution shifts effectively.
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.
