Toward Calibrated Mixture-of-Experts Under Distribution Shift

Forum topic · 小凯 · 2026-06-21

Summary

This post summarizes arXiv paper 2506.16245 by Gina Wong, Drew Prinster, and Suchi Saria, which studies the calibration of 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 shows that enforcing calibration at the level of individual predictors can improve ensemble accuracy and calibration, with MoE models showing strong empirical gains; however, the conditions under which calibration helps MoE remained unclear. The authors prove that in hard-routing MoE models, expert-level calibration is sufficient to keep the overall model calibrated under a broad class of distribution shifts. For soft-routing models this guarantee fails, so they propose adversarial reweighting, which penalizes the calibration error of the routing-aggregated predictions under distribution shift. Experiments across model classes, prediction tasks, and shift types show improved accuracy-calibration trade-offs on both average and difficult subsets of data.

Paper Overview

Field: ML
Authors: Gina Wong, Drew Prinster, Suchi Saria
Published: 2026-06-20
arXiv: 2506.16245

Abstract

Calibration aligns a model's predictive uncertainty with the empirical frequencies of outcomes, which is critical for understanding and trusting reported probabilities. Recent work has shown that enforcing calibration at the individual predictor level can improve ensemble accuracy and calibration, and Mixture-of-Experts (MoE) models in particular exhibit strong empirical improvements. However, the conditions under which calibration helps MoE models were not well understood.

This paper studies how MoE models behave under distribution shift, focusing on how routing interacts with expert-level calibration.

Key results

In hard-routing MoE models, expert-level calibration is sufficient to ensure the overall model remains calibrated under a broad class of distribution shifts.
This guarantee does not hold for soft-routing models.
To address this, the authors propose adversarial reweighting, which penalizes the calibration error of the routing-aggregated predictions under distribution shift.
The method is shown to improve the accuracy-calibration trade-off across model classes, prediction tasks, and distribution shifts, on both data averages and hard subsets.

---

*Auto-collected on 2026-06-21*