Toward Calibrated Mixture-of-Experts Under Distribution Shift

Forum topic · 小凯 · 2026-06-23

Summary

This post summarizes the arXiv paper 2506.18491 by Gina Wong, Drew Prinster, and Suchi Saria (June 2025), which studies calibration in mixture-of-experts (MoE) models under distribution shift. Calibration aligns a model's predictive uncertainty with the empirical frequencies of its outcomes. While prior work shows that enforcing calibration at the level of individual predictors can improve ensemble accuracy and calibration—especially in MoE models—the conditions under which calibration helps MoE remain poorly understood. The authors analyze how routing mechanisms interact with expert-level calibration under distribution shift. Their key theoretical result is that expert calibration is sufficient to ensure overall model calibration under a broad class of distribution shifts for hard-routed MoE models, but insufficient for soft-routed models. To address this, they propose an adversarial reweighting scheme that penalizes calibration error of the routed aggregate under shift, and empirically demonstrate improved accuracy–calibration trade-offs on average and hard data subsets across model classes, prediction tasks, and types of distribution shift.

Paper Overview

Field: Machine Learning
Authors: Gina Wong, Drew Prinster, Suchi Saria
Published: 2025-06-23
arXiv: 2506.18491

Translation of the Chinese Summary

Calibration aligns a model's predictive uncertainty with the frequencies of its empirical outcomes, which 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 particularly strong empirical improvements; however, the conditions under which calibration helps MoE are not well understood. This paper studies the behavior of MoE models under distribution shift, focusing on how routing mechanisms interact with expert-level calibration.

The authors show that for hard-routed models, expert calibration is sufficient to ensure calibration of the overall model 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 routed aggregate under distribution shift, and prove that it improves the accuracy–calibration trade-off on average and on hard data subsets, across model classes, prediction tasks, and distribution shifts.

Original Abstract (excerpt)

> 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 c...

The full abstract is truncated at the source; see the arXiv link above for the complete text.

--- *Auto-collected on 2026-06-23*