# Router Prior Bias: Preserving Base Routing Structure in MoE Post-Training

> Router Prior Bias, a soft logit anchoring to the frozen base router, outperforms load-balancing and unanchored fine-tuning in MoE post-training by preserving inherited routing structure.

- **Source:** [arXiv](https://arxiv.org/abs/2609.08115)
- **Published:** 2026-09-12
- **Permalink:** https://picx.dev/p/EScU92
- **Whiteboard:** https://picx.dev/p/EScU92/image

## Summary

## Summary (Overview)

- **Core contribution**: Introduces **Router Prior Bias (RPB)**, a training-time soft logit bias that anchors MoE routing to a prior derived from the frozen base router while keeping the router itself trainable, preserving inherited routing structure during post-training.
- **Key finding**: Re-applying the pretraining load-balancing loss (LBL) during post-training degrades downstream performance; soft anchoring outperforms both LBL and unanchored fine-tuning on Moonlight-16B-A3B (45.77 vs. 31.91 vs. 29.44 in-domain accuracy on math post-training).
- **Generality**: The advantage over LBL reproduces across model families (Qwen3-30B-A3B-Base) and corpora (OpenR1-Math-220k), while the advantage over no-load-balancing baseline is model-dependent.
- **Mechanistic insight**: The performance gain comes from the *softness* of the anchoring constraint, not from the specific prior content or the preservation of community structure—hard enforcement preserves community structure but sharply reduces performance.
- **Diagnostic protocol**: Introduces a three-level routing retention measurement (per-expert utilization, top-k overlap, community NMI) showing community structure is a *footprint* of soft anchoring rather than its source.

## Introduction and Theoretical Foundation

Mixture-of-Experts (MoE) language models route each token to a small subset of experts, allowing parameter growth without commensurate per-token computation. During pretraining, an **auxiliary load-balancing loss (LBL)** drives per-expert utilization toward uniformity (Lepikhin et al., 2020; Fedus et al., 2021). However, by post-training time, the router has already learned non-uniform specialization and recurring expert co-activation patterns (Tang et al., 2026; Lo et al., 2024). Re-imposing uniform utilization works *against* this acquired structure.

The paper identifies a gap in existing approaches:
- **Router-freeze** (strictest) removes trainability entirely, conflating *which* routing structure to preserve with *how strictly* to enforce it.
- **Unanchored fine-tuning** leaves base routing unconstrained.
- Practitioner toolchains (NVIDIA Megatron-Bridge, Unsloth) drop load balancing or disable router training, but without principled justification.

The authors separate two questions: **what routing structure to preserve** and **how strictly to enforce it**. RPB makes enforcement strength a continuous quantity ($\beta$) rather than an architectural commitment.

## Methodology

### Router Prior Bias (RPB)

**Frozen-base prior over Q/A spans**: For each training sample $x$, span $s \in \{Q, A\}$, layer $\ell$, and expert $e$, the frozen-base routing prior is:

$$
\pi_{x, s, \ell, e} = \mathrm{mean}_{t \in (x, s)} p_{\mathrm{base}}(e | x_t, \ell), \tag{2}
$$

where $p_{\mathrm{base}}$ is the frozen base router's softmax gate. Q/A spans are used because questions and answers induce different base routing distributions.

**Training-time bias**: With $c_{x,s,\ell,e} = \mathrm{clip}(\pi_{x,s,\ell,e})$ and $z_{t,\ell}$ the router-logit vector:

$$
b_{x, s, \ell, e} = \log c_{x, s, \ell, e} - \mathrm{mean}_{e'}\log c_{x, s, \ell, e'}, \tag{3}
$$

$$
z'_{t, \ell} = z_{t, \ell} + \beta b_{x, s, \ell}, \quad t \in (x, s). \tag{4}
$$

The clipping keeps log-bias finite ($\epsilon = 0.05$, $C = 5.0$), mean subtraction makes the bias a relative expert preference, and $\beta$ controls intervention strength ($\beta = 0$ recovers SFT).

### Comparison Methods

| Method | Constrained Object | Token-Conditioned Routing |
|---|---|---|
| **SFT / NoAux** | None | Preserved |
| **LBL** | Marginal expert load | Preserved |
| **RPB** | Router logits toward base prior | Preserved (soft bias) |
| **Router-L2** | Router weights toward base weights | Preserved |
| **KL-to-base** | Router distribution toward base | Preserved |
| **Logit-L2** | Centered router logits toward base | Preserved |
| **RPB-hard** | Base prior forces top-k | Removed |
| **within-soft** | Shuffled community prior, soft | Preserved |
| **within-hard** | Shuffled prior forces top-k | Removed |
| **global-hard** | Global shuffled prior forces top-k | Removed |

**Router-L2** adds a quadratic penalty:

$$
L = L_{\mathrm{SFT}} + \lambda \mathrm{mean}_{\ell}\left\| W_{\ell} - W_{\ell}^{\mathrm{base}} \right\|_2^2. \tag{5}
$$

### Routing Retention Metrics

Three levels are measured bias-free (training-time prior removed at evaluation):

1. **Per-expert utilization**: TVD and JSD between post-training and base expert distributions.
2. **Token-level top-k overlap**: Intersection size between base-router and post-training top-k sets, normalized by $k$.
3. **Expert co-activation community structure**: Layer-wise weighted expert graph with edge weights:

$$
w_{ij}^{(\ell)} = \frac{1}{N} \sum_{t=1}^{N} \mathbf{1}\{i, j \in \mathrm{top}_k(x_t)\}, \quad i \neq j, \tag{1}
$$

followed by Louvain community detection and NMI comparison.

## Empirical Validation / Results

### Main Results: Moonlight-16B-A3B

| Data | Method | ID | OOD |
|---|---|---|---|
| math20k | SFT | $29.44 \pm 0.32$ | $15.65 \pm 0.38$ |
| | LBL | $31.91 \pm 0.10$ | $14.97 \pm 0.51$ |
| | **RPB** | **$45.77 \pm 0.84$** | **$19.53 \pm 0.28$** |
| coding20k | SFT | $28.34 \pm 1.36$ | $35.83 \pm 0.49$ |
| | LBL | $30.73 \pm 0.13$ | $35.56 \pm 0.38$ |
| | **RPB** | **$41.37 \pm 0.30$** | **$38.70 \pm 0.11$** |

RPB attains the best mean on 21 of 24 (training-dataset, benchmark) cells.

### Anchoring Form Independence

Four anchoring objectives perform comparably (within ~1 point), regardless of whether they act on parameters, logits, or distributions:

| Method | math20k | coding20k |
|---|---|---|
| RPB (span-level logit bias) | $28.27 \pm 0.16$ | $39.59 \pm 0.14$ |
| Router-L2 (parameter space) | $27.69 \pm 0.28$ | $39.35 \pm 0.35$ |
| KL-to-base (probability space) | $27.50 \pm 0.62$ | $39.27 \pm 0.33$ |
| Logit-L2 (centered logit space) | $27.24 \pm 0.77$ | $39.81 \pm 0.67$ |
| NoAux / SFT | $20.25 \pm 0.29$ | $33.34 \pm 0.77$ |
| LBL (re-applied) | $20.61 \pm 0.33$ | $33.95 \pm 0.22$ |

### Generality

- **Qwen3-30B-A3B-Base**: LBL is weakest; RPB leads on math20k by 2.64 points (57.58 vs. 54.94), but the advantage over NoAux disappears.
- **OpenR1-Math-220k** (independent corpus): LBL costs 10.20 points against RPB (31.20 vs. 41.40), more than 20× the seed noise.

### Interventional Controls: Soft vs. Hard Enforcement

| Condition | math20k Overall@9 | NMI | coding20k Overall@9 | NMI |
|---|---|---|---|---|
| RPB (unshuffled) | 28.27 | 0.631 | 39.59 | 0.629 |
| within-soft | 27.23 | 0.632 | 38.96 | 0.632 |
| within-hard | 18.77 | 0.634 | 24.87 | 0.636 |
| global-hard | 3.59 | 0.642 | 17.96 | 0.650 |

Key observations:
1. **within-soft** stays within ~1 point of RPB → prior content need not name individual experts.
2. **within-hard** sharply reduces performance → enforcement strength is the operative axis.
3. **global-hard** reduces performance further → community corruption matters, yet its NMI is *highest*, showing NMI is a footprint, not an optimization target.

### Scope Condition: DeepSeek-V2-Lite

DeepSeek-V2-Lite's routing is close to uniform and nearly domain-invariant, so community NMI fails to separate interventions. Performance still improves under RPB, but the diagnostic only applies where base routers have sufficient community structure.

## Theoretical and Practical Implications

- **For MoE post-training**: The target should be preserving inherited co-activation structure rather than re-imposing uniform expert utilization. Flattening toward uniformity removes the property supporting out-of-domain retention.
- **For adaptation methods**: Methods that freeze or hard-select experts (ESFT, DES-MoE, MoCE) fix enforcement strength by construction; graded variants are worth testing against the range Router-L2 spans.
- **For evaluation methodology**: The dissociation between community NMI and performance warns that "distance from base model" proxies are informative only within their calibrated range—optimizing them directly pushes outside that range.
- **For architecture design**: What determines diagnostic applicability is how sharply the base router separates domains, not the expert architecture (shared-expert designs appear on both sides of the comparison).
- **Practical lesson**: Inherited routing should be held *softly* during post-training—both flattening it toward uniformity (LBL) and enforcing it absolutely (hard variants) carry downstream costs.

## Conclusion

Re-applying the pretraining load-balancing loss degrades downstream performance relative to soft router anchoring, with the effect surviving changes in model family and corpus. RPB is one competitive implementation of soft router anchoring rather than a uniquely necessary one—four objectives across different spaces perform comparably.

Two boundaries qualify the findings:
1. The advantage over a no-load-balancing baseline is model-dependent (reported both directions).
2. Community NMI is informative only where the base router carries community structure to begin with.

**Future directions**: The authors note that reasoning-trace post-training with long chains of thought is not covered. In such settings, the relevant object may be *sequential routing transition communities*—time-ordered patterns of expert-set transitions across reasoning steps—rather than static co-activation graphs. Constructing these transition graphs and re-running the shuffled-control protocol against them is a natural next step.

---

_Markdown view of https://picx.dev/p/EScU92, served by PicX — AI-generated visual whiteboard summaries of research papers._