Toward Calibrated Mixture-of-Experts Under Distribution Shift

Forum topic · 小凯 · 2026-06-22

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—aligning model-predicted uncertainties with empirical outcome frequencies—is essential for interpreting and trusting reported probabilities. Prior work has shown that enforcing calibration at the individual predictor level can improve both accuracy and calibration of ensemble models, with MoE architectures showing notable empirical gains, but the conditions under which calibration helps MoE remain unclear. The authors analyze the interaction between routing mechanisms and expert-level calibration, proving that expert calibration is sufficient to guarantee overall model calibration for hard-routed MoE under a broad class of distribution shifts, while it is insufficient for soft-routed models. They propose an adversarial reweighting method that penalizes calibration error of the routed aggregate under distribution shift, demonstrating improved accuracy-calibration trade-offs—both on average and on difficult data subsets—across model classes, prediction tasks, and distribution shift types. The paper falls under cs.AI and cs.LG and was released June 21, 2026.

Paper Overview

Research areas: cs.AI, cs.LG Authors: Gina Wong, Drew Prinster, Suchi Saria Published: 2026-06-21 arXiv: 2506.17580

Abstract

Calibration—aligning a model's predictive uncertainty with empirical outcome frequencies—is crucial for understanding and trusting reported probabilities. Recent work has shown that enforcing calibration at the level of individual predictors can improve the accuracy and calibration of ensemble models, with Mixture-of-Experts (MoE) models showing particularly notable empirical gains; however, the conditions under which calibration benefits MoE remain unclear.

This work studies the behavior of MoE models under distribution shift, focusing on the interaction between routing mechanisms and expert-level calibration. The authors prove that expert calibration is sufficient to ensure 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 routed aggregate under distribution shift. The method improves the accuracy-calibration trade-off—both on average and on difficult subsets of the data—across model classes, prediction tasks, and distribution shifts.

Key Findings

Hard routing: Expert-level calibration is sufficient to guarantee end-to-end model calibration under a wide class of distribution shifts.
Soft routing: Expert calibration alone does not ensure model-level calibration under shift.
Proposed method: An adversarial reweighting scheme that penalizes calibration error of the routed aggregate under shift, improving accuracy-calibration trade-offs across models, tasks, and shift types.

--- *Auto-collected on 2026-06-21*