Toward Calibrated Mixture-of-Experts Under Distribution Shift

Mixture-of-Experts (MoE) models are popular for their ability to scale compute efficiently by routing inputs to specialized experts. While it is often assumed that if individual experts are calibrated, the aggregate model will also be calibrated, this paper investigates whether this intuition holds under distribution shift. Specifically, the authors examine how different routing mechanisms (hard vs. soft) interact with expert-level calibration when the distribution of inputs—and thus the routing patterns—changes between training and deployment.

The authors analyze the conditions for aggregate calibration by treating the routing mechanism as a source of distribution shift. They distinguish between hard routing (where an input is assigned to one expert) and soft routing (where an input is a weighted combination of experts). By modeling the aggregate prediction as a function of the routing configuration, they show that soft routing creates a many-to-one mapping where different configurations can yield the same confidence score. They then introduce two training objectives, Robust MoE and Robust Filtered, which use entropy-balanced adversarial reweighting to penalize high-loss examples. This approach effectively forces the model to prioritize calibration on the most fragile routing configurations without requiring explicit knowledge of the test-time distribution.

The study reveals that expert-level calibration is sufficient to guarantee aggregate calibration under hard routing for a broad class of shifts. However, for soft-routed models, expert calibration is insufficient because aggregate confidence is determined by a delicate balance of routing weights and expert outputs. When a distribution shift alters the frequency of these configurations, the aggregate model becomes miscalibrated even if the experts themselves remain reliable. The proposed adversarial training objectives successfully mitigate this by focusing on high-loss examples, improving the accuracy-calibration tradeoff across various benchmarks without sacrificing overall accuracy.