Router prior bias improves mixture-of-experts post-training performance
Router Prior Bias: Preserving Base Routing Structure in MoE Post-Training
Artificial Intelligence
Summary
Mixture-of-experts models use multiple 'experts' to handle different parts of a problem, directed by a router that decides which expert to use. When fine-tuning these models after initial training, keeping the original router’s preferences with a gentle bias works better than forcing all experts to be used equally or changing the router freely. The authors show that softly holding onto the original routing patterns helps maintain performance and adaptability across tasks. They introduce a method called Router Prior Bias that gently guides the router toward its original behavior during post-training, improving accuracy and generalization.
What this means in practice
- •For machine learning engineers: Improve fine-tuning of large mixture-of-experts models by preserving routing patterns, boosting accuracy and robustness after pretraining.
- •For ai system developers: Maintain specialized routing behavior in expert models to enhance out-of-domain task performance during deployment adjustments.
Authors
Jaedeok Lee, Keonwoo Kim, Dongyoon Han, Sangdoo Yun, Yera Choi, Haanju Yoo
Abstract
Mixture-of-Experts (MoE) pretraining relies on an auxiliary load-balancing loss (LBL) to drive per-expert utilization toward uniformity. Post-training inherits a different situation: the base router already encodes non-uniform expert co-activation structure, which a re-imposed uniformity objective flattens away. We show that downstream performance depends instead on holding this inherited routing softly, a principle we term soft router anchoring, and instantiate it as Router Prior Bias (RPB), a training-time bias that pulls the router logits toward a prior read off the frozen base router while leaving the router itself trainable. On math post-training of Moonlight-16B-A3B, RPB attains 45.77 in-domain accuracy against 31.91 under re-applied LBL and 29.44 under unanchored fine-tuning, and retains more out-of-domain capability than either. The ordering against LBL reproduces on a second model family (Qwen3-30B-A3B-Base), and the advantage over LBL is resolvable on an independently sourced corpus. Anchors defined on the router weights, on its logits, or on its output distribution perform comparably with no consistent ordering, which places the effect in the softness of the constraint rather than in the particular prior RPB supplies. Retained community structure in the expert co-activation graph tracks these gains wherever the base router is non-uniform enough for communities to form, yet enforcing the same prior as a hard assignment preserves that structure while performance falls sharply. Community structure is therefore a footprint of soft anchoring rather than its source, and the practical lesson is that inherited routing should be held softly during post-training, since both flattening it toward uniformity and enforcing it absolutely carry a downstream cost. Our code will be released at https://github.com/naver-ai/rpb.