Toward Calibrated Mixture-of-Experts Under Distribution Shift

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.

• The paper introduces a robust calibration approach for MoE models under distribution shift using entropy-regularized adversarial objectives.
• It distinguishes calibration robustness between hard and soft routing, showing that soft routing is more prone to miscalibration due to aggregate prediction collapse.
• Empirical evaluations on CIFAR-10H, PACS, and CivilComments demonstrate significant reductions in calibration error with Robust MoE and Robust Filtered methods.

Calibration of Mixture-of-Experts Under Distribution Shift

Introduction

This paper presents an in-depth analysis of calibration properties in Mixture-of-Experts (MoE) models under distributional shift and proposes robust training objectives that directly address the routes by which miscalibration arises in soft-routed MoEs. The study provides a formally grounded distinction between the calibration robustness properties of hard- and soft-routing, characterizes the precise shift mechanisms that break calibration, and formulates entropy-regularized adversarial objectives for robust calibration. The empirical validation spans diverse domains, model backbones, and both synthetic and natural shifts, offering a rigorous evaluation of the proposed framework.

Theoretical Analysis

Calibration under Hard Routing

In the hard routing regime, each input is assigned to a single expert, forming a partition over the input space. The aggregate prediction coincides with that of the selected expert. Calibration of the global predictor reduces to calibration of the selected expert given its predicted probability: the sufficient statistic is the expert-identity and its output confidence. Calibration, in this setting, is robust to any input distributional shift that respects the conditional output distribution: as long as the label distribution in each expert-confidence "slice" is preserved, aggregate calibration holds.

Fragility of Soft-Routed MoEs

Soft routing, where the final prediction is a convex combination of expert outputs weighted by learned routing coefficients, breaks the bottleneck that confers robustness in hard routing. Here, the scalar probability output can result from exponentially many configurations of expert predictions and routing weights. Calibration at the expert level only enforces marginal constraints, failing to guarantee aggregate calibration due to the possibility of shifting the mix of configurations that collapse to the same scalar prediction. The main mathematical result is that the condition for invariance of calibration is substantially stronger—it requires that all configurations (routing vector + expert outputs) mapping to a particular prediction value have the same conditional label distribution. This is generically violated under shifts that change the routing distribution.

Robust Adversarial Training for Calibration

The core methodological contribution is a robust ERM framework where the MoE is trained against adversarial perturbations of the empirical distribution that upweight high-loss examples—a proxy for fragile routing configurations. The adversary is regularized via entropy balancing, ensuring that the reweighted distribution is close (in KL divergence) to the empirical distribution while maximizing expected loss. Two instantiations are provided:

• Robust MoE: Applies the entropy-tilted loss weighting to all examples.
• Robust Filtered: Targets only those samples likely to be routing-fragile (high disagreement among experts or large regret relative to the best expert).

This approach is connected to entropic DRO and multiaccuracy: the robust risk upper bounds the loss under bounded-KL reweightings of the input distribution and controls error across route-induced subpopulations.

Experimental Evaluation

The evaluation encompasses three regimes:

• CIFAR-10H (image classification with human-annotated ambiguities),
• PACS (domain-generalization benchmark with visually disparate domains),
• CivilComments (text toxicity classification under demographic subpopulation shift).

Baselines include vanilla ERM-trained MoEs, mixtures of individually calibrated experts (MoCaE), and frequency-aware gradient rectification (FGR). ECE and accuracy are evaluated in-distribution and on hard/shifted subpopulations.

Strong numerical results include:

• On CIFAR-10H low-agreement images, Robust MoE reduces Hard ECE from 0.281 (Vanilla MoE) and 0.262 (MoCaE) to 0.074, while FGR+Robust further improves this to 0.065.
• Robust Filtered yields the lowest overall ECE (0.013) on CIFAR-10H.
• On CivilComments, Robust MoE and Robust Filtered reduce Hard ECE from 0.108 (Vanilla MoE) and 0.101 (MoCaE) to 0.037 and 0.040, and maintain competitive accuracy.

These patterns are consistent across all metrics and datasets.

Figure 1: Reliability diagrams show non-robust MoEs place excessive probability mass in high-confidence bins, while robust objectives distribute probability more appropriately and bring predictions closer to empirical accuracies, confirming improved calibration.

Robust methods particularly excel under distribution shift—on ambiguous images, out-of-domain styles in PACS, and minority-mention text comments—where vanilla and even per-expert calibrated baselines exhibit overconfidence. Notably, post-hoc temperature scaling cannot recover the calibration deficit in MoEs under such shifts; only robust training at the configuration level is effective.

Subpopulation and Ambiguity Analysis

Results stratified by human agreement (CIFAR-10H) and demographic identity (CivilComments) affirm that robust training aligns confidence with accuracy even on ambiguous or underrepresented subsets. The mean confidence assigned by Robust MoE methods appropriately tracks the empirical accuracy as complexity or ambiguity increases, a behavior not mirrored by vanilla MoEs.

Implications and Future Directions

The results demonstrate that per-expert calibration and post-hoc recalibration are insufficient for aggregate reliability in MoEs due to the many-to-one collapse induced by the routing. The findings underscore the need for evaluating calibration at the aggregate level and under realistic or adversarial routing shifts. The strong performance of the robust objectives, especially Robust Filtered, suggests the value of targeting the sources of mixing-induced fragility, rather than only relying on global or marginal error indicators.

From a theoretical perspective, these findings motivate new formalizations of calibration and robustness for hierarchical or modular predictive architectures. For practice, they raise the bar for deploying MoE-based systems in safety- or fairness-critical contexts—evaluation and training under route-sensitive and subpopulation shifts is mandatory.

A pressing research direction is validating these phenomena at scale, particularly in sparse massive MoEs (as are becoming standard in LLM deployments), and extending adversarial configuration-aware calibration methods to generative modeling and structured prediction tasks.

Conclusion

The paper provides a rigorous and comprehensive study of calibration in mixture-of-experts under distribution shift, characterizing the explicit limitations of expert-level interventions and offering effective practical ameliorations by robust adversarial training. The implications span both foundational understanding of modular neural architectures and actionable strategies for deploying reliable AI at scale.