Gradient Rebalancing: Mitigating Exposure-Induced Routing Drift in MoE Fine-Tuning

Pretrained mixture-of-experts (MoE) models develop expert specialization across different input patterns. Narrow downstream distributions shift the relative frequencies of these patterns, naturally leading to imbalanced expert utilization under the pretrained router. Enforcing uniform utilization during fine-tuning may therefore disrupt useful pretrained routing preferences. However, simply removing load balancing leaves exposure asymmetry in router training unaddressed. Frequently selected experts receive direct gradient signals from more tokens. Their router updates can therefore have greater influence on routing boundaries than those of rarely selected experts. This imbalance creates a mismatch with the more balanced training conditions of pretraining, which may amplify deviations from pretrained routing behavior. To address this issue, we propose Gradient Rebalancing (GR) to mitigate exposure-induced routing drift without imposing load-balancing constraints. GR targets each expert's self-region tendency, a first-order estimate of the change in its routing-region mass induced by its own router update. We decompose this tendency into token-level contributions to region expansion and contraction. GR selectively attenuates these contributions to bring each expert's estimated tendency close to zero, regardless of its exposure. This limits exposure-driven asymmetry in boundary adjustments, helping reduce drift from pretrained routing behavior. By preserving the signs of individual gradient signals, GR retains flexibility for task-driven router adaptation. Experiments across diverse tasks and datasets show that GR achieves the best overall performance among the evaluated methods and substantially reduces routing drift with only a small computational overhead.