Router prior bias improves expert routing in large ai models
Router Prior Bias: Preserving Base Routing Structure in MoE Post-Training
Artificial Intelligence
Summary
Large AI models often use a system called Mixture-of-Experts (MoE) where different parts, called experts, handle different tasks. The way these experts are activated is controlled by a router, which is usually trained to spread work evenly across experts. The authors found that after initial training, forcing the router to act uniformly hurts performance. Instead, keeping the router's original patterns softly intact with a method called Router Prior Bias (RPB) helps maintain better accuracy and flexibility. This approach keeps the router adaptable while respecting its initial expert usage patterns, leading to better results.
Mixture-of-Experts (MoE)routerexpert co-activationload-balancing loss (LBL)soft router anchoringRouter Prior Bias (RPB)post-traininglogitsexpert community structure
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.