Rewiring Experts on the Fly: Continuous Rerouting for Better Online Adaptation in Mixture-of-Expert Models
Abstract
Mixture-of-Experts (MoE) models achieve efficient scaling through sparse expert activation, but often suffer from suboptimal routing decisions due to distribution shifts in deployment. While existing test-time adaptation methods could potentially address these issues, they primarily focus on dense models and require access to external data, limiting their practical applicability to MoE architectures. However, we find that, instead of relying on reference data, we can optimize MoE expert selection on-the-fly based only on input context. As such, we propose a data-free, online test-time framework that continuously adapts MoE routing decisions during text generation without external supervision or data. Our method cycles between two phases: During the prefill stage, and later in regular intervals, we optimize the routing decisions of the model using self-supervision based on the already generated sequence. Then, we generate text as normal, maintaining the modified router until the next adaption. We implement this through lightweight additive vectors that only update router logits in selected layers, maintaining computational efficiency while preventing over-adaptation. The results show consistent performance gains on challenging reasoning tasks while maintaining robustness to context shifts. For example, our method achieves a 5.5\% improvement on HumanEval with OLMoE. Furthermore, owing to its plug-and-play property, our method complements existing test-time scaling techniques, e.g., achieving 6\% average gains when incorporated with self-consistency on DeepSeek-V2-Lite.
Lay Summary
Large language models are increasingly built with Mixture-of-Experts architectures, where only a small subset of expert modules is activated for each input. This makes models more efficient, but it also creates a new problem: during deployment, the model may choose the wrong experts when the input distribution differs from what it saw during training. Existing test-time adaptation methods are not well suited to this setting, because they are mostly designed for dense models and often require extra reference data. We address this by adapting the expert-routing decisions directly during text generation, without using any external data or supervision. Our method uses the model’s own generated context to periodically improve which experts are selected, then continues generation with the updated routing behavior. This provides a practical way to make sparse expert models more robust at test time. By improving expert selection online, our work helps MoE models better handle real-world inputs where deployment conditions may differ from training data.