Rewiring Experts on the Fly:Continuous Rerouting for Better Online Adaptation in Mixture-of-Expert models

TL;DR

Mixture-of-Experts models suffer from suboptimal routing under deployment distribution shifts. The authors propose a data-free, online test-time rerouting method that continuously adapts MoE routing decisions by updating lightweight router logits in selected layers via self-supervised guidance from the generated context, alternating with steered generation. The approach uses phase-based optimization (In-Context Routing Optimization) and generation, with dynamic layer selection and additive delta parameters, achieving consistent improvements across reasoning benchmarks, including up to 6.7% on HumanEval and gains when combined with self-consistency. It is plug-and-play with existing test-time strategies and provides a practical pathway toward continual self-regulation of MoE models in real-world deployments.

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 \textit{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 experimental 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 naturally complements existing test-time scaling techniques, e.g., achieving 6\% average gains when incorporated with self-consistency on DeepSeek-V2-Lite.

Figures (7)

• Figure 1: Test-time rerouting framework for MoE models.(a) Rerouting mechanism: lightweight additive vectors (Delta) update router logits in selected high-confidence layers using self-supervised loss from existing context. (b) Continuous adaptation: alternating between optimization phases that adapt routing decisions and generation phases that maintain adapted routing until the next optimization cycle.
• Figure 2: Analysis of test-time rerouting mechanisms across different datasets in DeepSeek-V2-Lite. a) Edit distance before and after rerouting. b) Expert utilization dynamics before and after rerouting.
• Figure 3: Expert routing entropy as a function of sequence length averaged over 16 token blocks.
• Figure 4: Performance of our method versus the baseline across different few-shot examples under shifted and aligned task contexts.
• Figure 5: Performance of our method versus the baseline on AIME datasets using the GPT-OSS-20b model. Optimizing rerouting is especially noticeable at improving model certainty, as shown by improved performance at lower pass@k and improved majority voting, even one challenging benchmarks like AIME.