Overcoming Data Imbalances with Expert Selection Biases

Ayush K. Tarun; Clara Meister; Antoine Bosselut

cs.CLOct 9, 2026

Overcoming Data Imbalances with Expert Selection Biases

Authors: Ayush K. Tarun, Clara Meister, Antoine Bosselut

Organizations: EPFL

Abstract

Mixture-of-Experts (MoE) models have become the standard architecture for frontier LLMs. The multiple parallel experts in each transformer block theoretically enable expert specialization according to categories such as domains, languages, or modalities during pretraining. However, when data is imbalanced across categories, minority categories can become concentrated in a small fraction of the expert pool or spread inefficiently across them, effectively limiting their available capacity. In this work, we demonstrate that category-conditioned expert-selection biases can preferentially route tokens from different categories toward specialized expert pools. By applying biases based on input category, we provide a simple way to induce soft category-aligned specialization. This intervention encourages categories toward distinct expert subsets while retaining cross-category sharing when useful. More importantly, it enables the over-allocation of expert capacity to categories that would otherwise be disadvantaged by data scarcity. We observe performance improvements across modalities, low-resource languages, and code in models up to 10.6B parameters. At our smallest and largest model scales, these gains are equivalent to allocating 18.5% and 2.7% additional baseline training tokens, respectively. Overall, our results show that expert biases can serve as a lightweight control for category-level expert allocation, making room for low-resource categories while retaining standard MoE inference.

Figures & tables

Figure 1: Given an input category, COAX adds a category-conditioned bias

\alpha

to the routing scores of a designated subset of experts before top-

k

selection. This bias softly encourages tokens from the same category to route to a common expert set, promoting category-specific specialization, while leaving all experts available when the learned routing scores sufficiently favor them. The mechanism is applied across MoE layers during pretraining, but the biases can be removed at inference time.

Figure 2: COAX effectively coaxes category-aligned routing. In-group routing fraction: the fraction of a token’s top-

K

selected experts that belong to its designated subset

\mathcal{P}_{c_{t}}

. Dashed lines show the uniform-random routing rate

q_{c}

which would be the expected fraction for an unbiased baseline. Larger bias strengths (

\alpha

) consistently produce stronger routing toward the designated expert subset.

Model	Inf. Bias	English	Korean	Urdu	Swahili
345M-A147M	No	$+0.05$	$\mathbf{-0.01}$	$\mathbf{-0.23}$	$\mathbf{-0.55}^{**}$
345M-A147M	Yes	$\mathbf{-0.02}$	$\mathbf{-0.08}$	$\mathbf{-0.28}$	$\mathbf{-0.59}^{**}$
3B-A654M	No	$\mathbf{-0.01}$	$\mathbf{-0.11}$	$\mathbf{-0.31}^{*}$	$\mathbf{-0.40}^{**}$
3B-A654M	Yes	$\mathbf{-0.09}^{*}$	$\mathbf{-0.18}$	$\mathbf{-0.38}^{*}$	$\mathbf{-0.44}^{**}$
10.6B-A1.2B	No	$+0.02$	$\mathbf{-0.00}$	$\mathbf{-0.10}^{*}$	$\mathbf{-0.30}^{**}$
10.6B-A1.2B	Yes	$\mathbf{-0.09}$	$\mathbf{-0.13}^{*}$	$\mathbf{-0.22}^{**}$	$\mathbf{-0.41}^{**}$

Table 1: Mean relative change in validation loss (

\Delta\%

) against unbiased baselines trained with matched seeds (negative is better). ’Inf. Bias’ denotes whether the category-conditioned bias is retained during inference. Bold indicates improvements over the unbiased baseline (

\Delta<0

). ∗ and ∗∗ denote paired differences whose

90\%

and

95\%

confidence intervals respectively exclude zero. Our results demonstrate that COAX more significantly helps lower-resource languages with minimal impact to higher-resource languages.

Language	COAX	Hard-routing	EDL	Language	COAX	Hard-routing	EDL
English	$\mathbf{-0.070}$	$+0.350$	$+0.047$	Tamil	$-0.418$	$\mathbf{-0.423}$	$+0.400$
Chinese	$\mathbf{-0.103}$	$+1.994$	$+0.222$	Georgian	$\mathbf{-0.576}$	$+0.228$	$+1.257$
Korean	$\mathbf{-0.104}$	$+1.679$	$+0.091$	Urdu	$\mathbf{-0.308}$	$-0.101$	$+0.926$
Khmer	$-0.003$	$\mathbf{-4.090}$	$+1.193$	Swahili	$\mathbf{-0.486}$	$+1.039$	$+3.020$

Table 2: Relative change in validation loss (

\Delta\%

) for the 10.6B-A1.2B model trained on eight languages against an unbiased baseline (negative is better). Columns correspond to our proposed method COAX, DEMix-style hard routing (Hard-routing) and Expert Divergence Learning (EDL). Bold denotes the best result for each language. Hard-routing and EDL induce category-aligned specialization without consistent performance benefits in this setting, highlighting the importance of COAX’s ability to allocate different preferred expert shares across categories.

Model	Inf. Bias	T	I2T	T2I	AA	AT	ATA
345M-A147M	No	$\mathbf{-0.78}^{*}$	$\mathbf{-0.72}^{**}$	$\mathbf{-0.56}^{**}$	$\mathbf{-0.48}$	$\mathbf{-0.38}$	$\mathbf{-0.65}$
345M-A147M	Yes	$\mathbf{-0.88}^{*}$	$\mathbf{-0.76}^{**}$	$\mathbf{-0.59}^{**}$	$\mathbf{-0.50}^{*}$	$\mathbf{-0.39}$	$\mathbf{-0.67}$
3B-A654M	No	$+0.23^{**}$	$\mathbf{-0.25}^{**}$	$\mathbf{-0.19}^{**}$	$+0.04$	$+0.11$	$+0.09$
3B-A654M	Yes	$+0.18^{**}$	$\mathbf{-0.29}^{**}$	$\mathbf{-0.22}^{**}$	$+0.02$	$+0.09$	$+0.08$
10.6B-A1.2B	No	$\mathbf{-0.05}$	$\mathbf{-0.25}^{*}$	$\mathbf{-0.18}^{*}$	$\mathbf{-0.08}$	$\mathbf{-0.09}$	$\mathbf{-0.09}$
10.6B-A1.2B	Yes	$\mathbf{-0.13}$	$\mathbf{-0.31}^{**}$	$\mathbf{-0.23}^{**}$	$\mathbf{-0.12}^{*}$	$\mathbf{-0.12}^{*}$	$\mathbf{-0.13}$

Table 3: Mean relative change in validation loss (

\Delta\%

) against unbiased baselines trained with matched seeds (negative is better). Columns correspond to held-out text (T), image-to-text captioning (I2T), text-to-image generation (T2I), audio-only (AA), audio-to-text ASR (AT), and interleaved audio-text (ATA) subtasks in validation data. Vision and audio are each biased toward approximately

2\times

their token share, with text receiving the residual expert subset (

r_{c}\approx 0.7

). Inf. Bias denotes whether the category-conditioned bias is retained during inference. Bold indicates improvements over the unbiased baseline (

\Delta<0

). ∗ and ∗∗ denote paired differences whose

90\%

and

95\%

confidence intervals respectively exclude zero. Our results show that COAX improves data-scarce modalities with minimal impact on the higher-resource text modality.

Model	Inf. Bias	English	Code	Math	Legal	Medical
345M-A147M	No	$+0.26$	$\mathbf{-0.19}$	$+0.10$	$+0.04$	$+0.07$
345M-A147M	Yes	$+0.14$	$\mathbf{-0.35}$	$\mathbf{-0.05}$	$\mathbf{-0.15}$	$\mathbf{-0.11}$
3B-A654M	No	$+0.20^{**}$	$\mathbf{-0.10}$	$+0.62$	$+0.05$	$+0.09$
3B-A654M	Yes	$+0.10$	$\mathbf{-0.25}$	$\mathbf{-0.12}$	$\mathbf{-0.09}$	$\mathbf{-0.06}$
10.6B-A1.2B	No	$+0.25^{**}$	$+0.00$	$+0.10$	$+0.05$	$+0.11$
10.6B-A1.2B	Yes	$+0.07$	$\mathbf{-0.19}^{*}$	$\mathbf{-0.04}$	$\mathbf{-0.08}$	$\mathbf{-0.04}$

Table 4: Mean relative change in validation loss (

\Delta\%

) against matched-seed unbiased baselines (negative is better) . English is under-provisioned (

r_{c}\approx 0.7

), code and mathematics use

r_{c}=2

, and legal and medical use

r_{c}=4

. Inf. bias denotes whether the category-conditioned bias is retained during inference. Bold denotes improvement over the baseline; ∗ and ∗∗ denote paired differences whose

90\%

and

95\%

confidence intervals respectively exclude zero. Our results demonstrate that COAX may help specialized domains when the expert bias is kept at inference time, while removing it only helps code.

Language	$\bm{\Delta}$ loss (%)		Language	$\bm{\Delta}$ loss (%)
	Per-lang.	Eng. vs rest		Per-lang.	Eng. vs rest
English	$-0.07$	$\mathbf{-0.15}$	Tamil	$-0.42$	$\mathbf{-0.53}$
Chinese	$-0.10$	$\mathbf{-0.20}$	Georgian	$-0.58$	$-0.58$
Korean	$-0.10$	$\mathbf{-0.29}$	Urdu	$-0.31$	$\mathbf{-0.37}$
Khmer	$\mathbf{0.00}$	$+0.19$	Swahili	$\mathbf{-0.49}$	$-0.26$

Table 5: Fine-grained expert partitioning is not necessary. Relative change in validation loss for the 10.6B-A1.2B model trained on the eight-language mixture from Section 5.1 ; negative is better. Per-lang. assigns each language its own expert subset, whereas Eng. vs rest uses only two groups: English and all non-English languages. Bold denotes the better grouping for each language. The coarser split performs as well as or better than per-language partitioning for six of eight languages, suggesting that much of COAX’s benefit comes from separating the dominant language from the long tail rather than separating every minority language from each other.

	No inference bias ( $\bm{\Delta}$ loss % )		With inference bias ( $\bm{\Delta}$ loss % )
$\bm{\alpha}$	English	Minority mean	English	Minority mean	$\bm{\Delta}$ TFLOP/s
$0.003$	$-0.02$	$-0.12$	$-0.02$	$-0.12$	$0.0\%$
$0.01$	$-0.10$	$-0.30$	$\mathbf{-0.11}$	$\mathbf{-0.31}$	$-3.9\%$
$0.03$	$-0.03$	$-0.42$	$\mathbf{-0.12}$	$\mathbf{-0.48}$	$-13.1\%$
$0.1$	$+0.38$	$-0.37$	$\mathbf{-0.04}$	$\mathbf{-0.96}$	$-27.3\%$
$0.3$	$+1.88$	$+19.32$	$\mathbf{+0.46}$	$\mathbf{-1.65}$	$-25.6\%$

Table 6: Effect of bias strength. Relative change in validation loss for English and the mean over the three minority languages: Korean, Urdu, and Swahili, for the 3B-A654M model on the four-language mixture from Section 5.1 ; negative is better. No inference bias removes the category bias at inference, whereas with inference bias retains the training-time bias. Without validation-time bias, gains peak around

\alpha=0.03

before train-test mismatch dominates at larger values. Retaining the bias yields increasingly large minority-language gains as

\alpha

increases, at the cost of greater inference overhead. In the final row,

\alpha=0.3

is linearly annealed to zero during cooldown, recovering

76.4\%

of the corresponding minority-language gain without requiring biasing at inference.

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Hidden dim.	Expert dim.	# Experts	# Heads	# KV Heads
345M-A147M	384	96	64	16	4
3B-A654M	1280	160	128	20	5
10.6B-A1.2B	2048	192	256	16	4

Appendix

Table 7: Architecture configurations of the MoE models used in our experiments. All models use 32 layers, one shared expert, and activate

K=8

routed experts per token.

Language	$n$	Base	COAX	$\Delta$	Base	COAX	$\Delta$
		Unnormalised			Length-normalised
Tamil	$382$	$98.17$	$97.91$	$-0.26$	$90.05$	$91.62$	$+1.57$
Georgian	$204$	$96.57$	$96.08$	$-0.49$	$91.67$	$92.16$	$+0.49$
Urdu	$550$	$96.55$	$96.91$	$+0.36$	$95.45$	$95.09$	$-0.36$

Appendix

Table 8: MultiBLiMP shows no consistent COAX effect at our training scale. Accuracy (%) for the 10.6B-A1.2B eight-language models evaluated without inference-time bias; higher is better. Differences are small and inconsistent across the two scoring schemes, while both models retain high absolute performance.

		Token-level			Word-level
Task		top-1	top-5	top-10	top-1	top-5	top-10
ASR ( $n{=}2048$ )	Base	$32.10$	$55.63$	$64.24$	$24.41$	$47.50$	$57.18$
	COAX	$\mathbf{33.84}$	$\mathbf{57.29}$	$\mathbf{65.75}$	$\mathbf{26.07}$	$\mathbf{49.38}$	$\mathbf{58.98}$
	$\Delta$	$+1.74$	$+1.66$	$+1.51$	$+1.66$	$+1.88$	$+1.80$
Captioning ( $n{=}1024$ )	Base	$57.33$	$81.98$	$87.63$	$50.44$	$78.52$	$85.20$
	COAX	$\mathbf{58.34}$	$\mathbf{82.89}$	$\mathbf{88.29}$	$\mathbf{51.49}$	$\mathbf{79.57}$	$\mathbf{85.99}$

Appendix

Table 9: COAX improves task-aligned ASR and captioning predictions. Teacher-forced top-

k

accuracy (%) on reference transcripts and captions for a 345M-A147M model, evaluated without inference-time bias; higher is better. A word is counted as correct at

k

only when all of its constituent tokens appear within the top-

k

predictions. COAX improves all twelve metrics.

New category	New-category $\Delta$ loss (%)	Replay $\Delta$ loss (%)
Vision	$\mathbf{-0.066}$	$\mathbf{-0.067}$
Code	$\mathbf{-0.177}$	$\mathbf{-0.064}$
Russian	$+0.028$	$\mathbf{-0.054}$
Hindi	$+0.029$	$\mathbf{-0.039}$

Appendix

Table 10: COAX can allocate newly added expert capacity during continual pretraining. Relative change in validation loss (

\Delta\%

) against unbiased continual pretraining from the same checkpoint for the 3B-A654M model (from Section 5.1 ) upcycled from

128

160

routed experts at

\alpha=0.03

; negative is better. The new category receives

32/160

experts and

20\%

of training tokens. Replay is the mean over the four original languages.

Figure 3: COAX does not sharpen the learned router. Router entropy over all

128

experts, averaged across tokens and layers, for the 3B-A654M model on the four-language mixture from Section 5.1 . Entropy is computed before adding the category bias and therefore reflects the router’s learned scores. Both models sharpen during training, but the COAX-trained router remains consistently flatter, suggesting that its performance gains arise from changes in expert learning rather than a better, more discriminate router.

Disabled set	Model	Swahili traffic (%)	Swahili	English	Korean	Urdu
Swahili-targeted	Baseline	$4.58$	$+9.34$	$+0.92$	$+0.71$	$+0.97$
	COAX $\alpha{=}0.03$	$12.15$	$+107.09$	$+0.75$	$+0.83$	$+1.59$
	COAX $\alpha{=}0.3$ annealed	$\mathbf{26.96}$	$\mathbf{+128.18}$	$\mathbf{+0.08}$	$\mathbf{+0.19}$	$\mathbf{+0.03}$
Arbitrary	Baseline	$2.99$	$+0.87$	$+0.93$	$+0.76$	$+0.68$
	COAX $\alpha{=}0.03$	$2.70$	$+0.54$	$+0.87$	$+0.86$	$+0.47$
	COAX $\alpha{=}0.3$ annealed	$2.10$	$+0.25$	$+0.78$	$+0.90$	$+0.04$

Appendix

Table 11: COAX concentrates Swahili-specific functionality into its designated experts. Relative change in validation loss (

\Delta\%

) after disabling four of

128

experts at every layer of the 3B-A654M model; larger positive values indicate greater functional dependence. Swahili traffic is the fraction of Swahili top-

K

selections routed through the disabled experts in the intact model; the uniform-random rate is

3.12\%

(

4/128

). All models are evaluated without inference-time bias.

Language	Experts	COAX ( $\alpha=0.03$ )	Hard, weights-only	Hard + bias on eval
English	$114$	$-0.070$	$+15.868$	$+0.350$
Chinese	$26$	$-0.103$	$+88.817$	$+1.994$
Korean	$26$	$-0.104$	$+118.728$	$+1.679$
Khmer	$30$	$-0.003$	$+604.163$	$\mathbf{-4.090}$
Tamil	$22$	$-0.418$	$+253.411$	$\mathbf{-0.423}$
Georgian	$16$	$-0.576$	$+358.344$	$+0.228$

Appendix

Table 12: COAX outperforms generalized DEMix-style hard routing. Relative change in validation loss (

\Delta\%

) against the same-seed unbiased baseline for the 10.6B-A1.2B eight-language setup from Section 5.1 ; negative is better. Hard, weights-only removes the routing bias at evaluation, whereas hard + bias on eval retains it. Iteration-time overhead is measured using the same throughput probe as Appendix G.1 .

Language	Share	COAX ( $\alpha=0.03$ )	EDL ( $\beta=5\!\times\!10^{-4}$ )	EDL ( $\beta=10^{-4}$ )	EDL ( $\beta=2\!\times\!10^{-5}$ )
English	$63.2\%$	$\mathbf{-0.070}$	$+0.047$	$+0.082$	$+0.065$
Chinese	$10.0\%$	$\mathbf{-0.103}$	$+0.222$	$+0.028$	$+0.098$
Korean	$10.0\%$	$\mathbf{-0.104}$	$+0.091$	$-0.013$	$+0.063$
Khmer	$6.0\%$	$\mathbf{-0.003}$	$+1.193$	$+3.866$	$+0.738$
Tamil	$4.2\%$	$\mathbf{-0.418}$	$+0.400$	$+0.319$	$+0.336$
Georgian	$3.0\%$	$\mathbf{-0.576}$	$+1.257$	$+0.523$	$+0.502$

Appendix

Table 13: EDL does not improve the long-tailed multilingual mixture at any of the tested coefficients. Relative change in validation loss (

\Delta\%

) against the matched unbiased baseline for the 10.6B-A1.2B eight-language setup of Section 5.1 ; negative is better. Share denotes each language’s fraction of the training mixture.

Method	Aux. loss	Excess over uniform	$\mathbf{\mathrm{MaxVio}_{\mathrm{batch}}}$	Router JSD
Unbiased baseline	$1.0345$	$3.45\%$	$0.4184$	$0.0025$
COAX ( $\alpha=0.03$ )	$1.0346$	$3.46\%$	$0.6568$	$0.0027$
EDL ( $\beta=2\!\times\!10^{-5}$ )	$1.0581$	$5.81\%$	$0.4618$	$0.0296$
EDL ( $\beta=10^{-4}$ )	$1.1063$	$10.63\%$	$0.5606$	$0.1406$
EDL ( $\beta=5\!\times\!10^{-4}$ )	$1.2335$	$23.35\%$	$0.4671$	$0.3300$

Appendix

Table 14: EDL strongly separates language-level router distributions, while COAX leaves the probability-level divergence nearly unchanged. Statistics are averaged over the final

100

training steps, except Router JSD , which is measured on the final checkpoint. The auxiliary loss is normalized such that uniform routing gives

1

\mathrm{MaxVio}_{\mathrm{batch}}

( Wang et al., 2024 ) measures the relative excess load of the most-loaded expert over its load under perfect balance. Router JSD is the mean pairwise Jensen-Shannon divergence between language-level router distributions, with maximum

\ln 2

Model	Method	T	I2T	T2I	AA	AT	ATA
345M-A147M	$r_{c}\!=\!1$	$-0.28$	$-0.30$	$-0.22$	$-0.23$	$-0.20$	$-0.35$
	+ bias at eval	$-0.41$	$-0.34$	$-0.25$	$-0.25$	$-0.21$	$-0.37$
	$r_{c}=2$	$\mathbf{-0.78}^{*}$	$\mathbf{-0.72}^{**}$	$\mathbf{-0.56}^{**}$	$\mathbf{-0.48}$	$\mathbf{-0.38}$	$\mathbf{-0.65}$
	+ bias at eval	$\mathbf{-0.88}^{*}$	$\mathbf{-0.76}^{**}$	$\mathbf{-0.59}^{**}$	$\mathbf{-0.50}^{*}$	$\mathbf{-0.39}$	$\mathbf{-0.67}$
3B-A654M	$r_{c}\!=\!1$	$+0.47^{*}$	$-0.17^{**}$	$-0.14^{**}$	$+0.13$	$+0.16^{**}$	$+0.18^{*}$
	+ bias at eval	$+0.39^{*}$	$-0.20^{**}$	$-0.17^{**}$	$+0.11$	$+0.15^{**}$	$+0.17^{*}$

Appendix

Table 15: COAX helps multimodal performance. Mean relative change in validation loss (

\Delta\%

) against unbiased baselines trained with matched seeds; negative is better. Columns correspond to held-out text (T), image-to-text captioning (I2T), text-to-image generation (T2I), audio-only (AA), audio-to-text ASR (AT), and interleaved audio-text (ATA). For

r_{c}=1

, each modality is assigned an expert fraction matching its token share. For

r_{c}=2

, vision and audio are each biased toward approximately

2\times

their token share, with text receiving the residual expert subset (

r_{c}\approx 0.7

). ’+ bias at eval’ retains the category-conditioned bias during inference; otherwise, it is removed. ∗ and ∗∗ denote paired differences whose

90\%

and

95\%

confidence intervals respectively exclude zero. Within each evaluation setting, bold identifies the better of

r_{c}\in\{1,2\}

when COAX improves over the unbiased baseline (

\Delta<0

Swahili experts	Swahili	English	Korean	Urdu
$2$	$+0.043\pm 0.249$	$-0.026\pm 0.044$	$-0.001\pm 0.223$	$+0.088\pm 0.347$
$4$	reference
$8$	$-0.087\pm 0.158$	$-0.032\pm 0.062$	$-0.070\pm 0.100$	$-0.027\pm 0.143$
$16$	$+0.138\pm 0.119$	$-0.002\pm 0.029$	$-0.021\pm 0.131$	$+0.049\pm 0.253$
$32$	$+0.207\pm 0.215$	$+0.010\pm 0.045$	$+0.061\pm 0.080$	$+0.089\pm 0.194$

Appendix

Table 16: COAX is robust to the exact relative expert share. Varying only the number of experts assigned to Swahili while keeping Korean and Urdu fixed for the 3B-A654M four-language setup from Section 5.1 at

\alpha=0.03

. Values report

\Delta\%

relative to the

4

-expert configuration at the same seed, averaged over three seeds

\pm

standard deviation; negative is better.

Biased layers	English	Korean	Urdu	Swahili	Mean
All ( $1$ – $32$ )	$+0.041$	$+0.031$	$\mathbf{-0.142}$	$\mathbf{-0.563}$	$\mathbf{-0.158}$
Early ( $1$ – $21$ )	$+0.134$	$+0.215$	$+0.164$	$\mathbf{-0.093}$	$+0.105$
Middle ( $12$ – $21$ )	$+0.346$	$+0.465$	$+0.405$	$+0.310$	$+0.382$
Late ( $22$ – $32$ )	$+0.092$	$+0.069$	$+0.056$	$\mathbf{-0.333}$	$\mathbf{-0.029}$

Appendix

Table 17: Restricting COAX to specific layer ranges does not improve over biasing all layers. Relative change in validation loss (

\Delta\%

) against matched-seed unbiased baselines for the 345M-A147M four-language setup at

\alpha=0.03

, averaged over three seeds 28–30; negative is better. The mean is over all four languages.

Bias strength	Aux coeff.	English	Korean	Urdu	Swahili	Minority mean
$\alpha=0.03$	$1e^{-4}$	$+0.17$	$\mathbf{-0.07}$	$\mathbf{-0.14}$	$\mathbf{-0.33}$	$\mathbf{-0.18}$
	$1e^{-3}$	$\mathbf{-0.03}$	$\mathbf{-0.26}$	$\mathbf{-0.59}$	$\mathbf{-0.41}$	$\mathbf{-0.42}$
	$1e^{-2}$	$+0.03$	$\mathbf{-0.06}$	$+0.02$	$\mathbf{-0.51}$	$\mathbf{-0.18}$
$\alpha=0.3$ annealed	$10^{-4}$	$+0.48$	$\mathbf{-0.40}$	$\mathbf{-1.69}$	$\mathbf{-1.68}$	$\mathbf{-1.26}$
	$10^{-3}$	$+0.57$	$\mathbf{-0.49}$	$\mathbf{-1.82}$	$\mathbf{-1.48}$	$\mathbf{-1.26}$
	$10^{-2}$	$+0.67$	$+0.04$	$\mathbf{-0.41}$	$\mathbf{-0.91}$	$\mathbf{-0.43}$

Appendix

Table 18: COAX remains effective under standard load balancing, but strong load-balancing suppresses it. Relative change in validation loss (

\Delta\%

) against an unbiased baseline trained with the same auxiliary coefficient for the 3B-A654M four-language setup from Section 5.1 ; negative is better. The minority mean averages Korean, Urdu, and Swahili. All evaluations remove the inference-time bias.

Model	Baseline	$r_{c}=1$	$r_{c}=2$
345M-A147M	$1.0312$	$+0.31\%$	$+0.26\%$
3B-A654M	$1.0202$	$-0.08\%$	$-0.24\%$
10.6B-A1.2B	$1.0305$	$-0.06\%$	$-0.22\%$

Appendix

Table 19: COAX barely changes the auxiliary load-balancing loss. Statistics are averaged over the final

100

training steps of the multimodal runs in Table 3 . Baseline values are absolute auxiliary losses, while COAX columns report relative changes against the matched baseline, averaged over five seeds at 345M-A147M and four seeds at the larger scales.

Figure 4: COAX’s gain does not decay over training steps. Relative reduction in validation loss against the unbiased baseline at intermediate checkpoints of the 10.6B-A1.2B two-group run (Section 5 ); positive favors COAX. After an initially larger gain, the advantage stabilizes from approximately

5

B through the final

10

B-token checkpoint.

Figure 5: COAX’s advantage persists through 100B training tokens. Relative reduction in validation loss against the unbiased baseline for the 10.6B-A1.2B two-group multilingual run; positive favors COAX. After increasing slightly from 10B to 20B tokens, the gain remains stable through the final 100B-token checkpoint.

Language	COAX	No English bias	Difference
English	$-0.070$	$-0.090$	$-0.020$
Chinese	$-0.103$	$-0.003$	$+0.100$
Korean	$-0.104$	$+0.030$	$+0.135$
Khmer	$-0.003$	$+0.729$	$+0.733$
Tamil	$-0.418$	$-0.132$	$+0.288$
Georgian	$-0.576$	$-0.211$	$+0.367$

Appendix

Table 20: Biasing the majority category is important. Relative change in validation loss (

\Delta\%

) against the same-seed unbiased baseline for the 10.6B-A1.2B eight-language setup from Section 5.1 at

\alpha=0.03

; negative is better. No English bias retains the same expert allocation and minority biases as COAX but removes the routing bias for English. The final column is No English bias minus COAX , so positive values favor full COAX.

Language	COAX: English gets $44.5\%$ of experts	English reduced to $44.5\%$ of data
English	$\mathbf{-0.070}$	$+6.583$
Chinese	$\mathbf{-0.103}$	$+2.306$
Korean	$\mathbf{-0.104}$	$+1.822$
Khmer	$\mathbf{-0.003}$	$+1.045$
Tamil	$\mathbf{-0.418}$	$+1.864$
Georgian	$\mathbf{-0.576}$	$+2.554$

Appendix

Table 21: Reducing majority data and reducing majority expert capacity are not the same. Relative change in validation loss (

\Delta\%

) against the same-seed unbiased baseline for the 10.6B-A1.2B eight-language setup from Section 5.1 ; negative is better. In the data-reduction run, every minority retains its original absolute token count while only English is reduced. Reducing English data regresses performance across all languages, unlike COAX.

Arm	English	Korean	Urdu	Swahili	Minority mean
COAX only	$+0.011$	$\mathbf{-0.063}$	$\mathbf{-0.364}$	$\mathbf{-0.352}$	$\mathbf{-0.260}$
Data only ( $2\times$ tail)	$+0.276$	$+0.273$	$\mathbf{-5.317}$	$\mathbf{-5.694}$	$\mathbf{-3.579}$
Both	$+0.274$	$+0.126$	$\mathbf{-5.420}$	$\mathbf{-5.921}$	$\mathbf{-3.738}$

Appendix

Table 22: COAX and data upsampling are additive. Relative change in validation loss (

\Delta\%

) against matched-seed unbiased baselines for the 3B-A654M four-language setup at

\alpha=0.03

, averaged over three seeds; negative is better. The data run doubles Urdu and Swahili to

2

epochs at

+3.61\%

training compute. The minority mean averages Korean, Urdu, and Swahili.

	Multimodality		Multilinguality
Model	$\Delta$ TFLOP/s	$\Delta$ time/iter	$\Delta$ TFLOP/s	$\Delta$ time/iter
345M-A147M	$-17.0\%$	$+20.2\%$	$-17.7\%$	$+21.9\%$
3B-A654M	$-11.6\%$	$+13.2\%$	$-13.2\%$	$+15.1\%$
10.6B-A1.2B	$-2.4\%$	$+2.4\%$	$-1.7\%$	$+1.7\%$
10.6B-A1.2B (Coarser)	-	-	$-1.3\%$	$+1.3\%$

Appendix

Table 23: Training throughput overhead of COAX at

\alpha=0.03

, measured over steps 8050–8200 from the step-8000 checkpoint. Negative

\Delta

TFLOP/s and positive

\Delta

time/iter indicate slower training relative to the unbiased baseline. The overhead decreases sharply with model scale, reaching only

1.7

–

2.4\%

at 10.6B-A1.2B. ’Coarser’ denotes the coarser configuration of COAX from Section 5 .

Language	Data %	COAX benefit	Equivalent extra training
English	$63.24$	$-0.170$	$+1.7\%$
Chinese	$10.00$	$-0.191$	$+2.9\%$
Korean	$10.00$	$-0.275$	$+3.7\%$
Khmer	$6.00$	$+0.254$	—
Tamil	$4.19$	$-0.494$	$+6.3\%$
Georgian	$2.99$	$-0.527$	$+5.8\%$

Appendix

Table 24: At 10.6B, COAX is worth

2.7\%

additional baseline training. Results for the 10.6B-A1.2B two-group eight-language configuration at

\alpha=0.03

, evaluated without inference-time bias. COAX benefit is the relative change in validation loss (

\Delta\%

) at step

8000

(no cooldown, unlike the final checkpoint numbers from Section 5.1 ); negative favors COAX. Equivalent extra training is the additional number of baseline training tokens required to reach the corresponding COAX loss, obtained by continuing the baseline at the same learning rate and interpolating between evaluated checkpoints. Khmer has no equivalence point because COAX is worse on it at the reference checkpoint.

	T	I2T	T2I	AA	AT	ATA
COAX benefit at $10$ B	$-0.613$	$-0.531$	$-0.399$	$-0.498$	$-0.206$	$-0.471$
Baseline gap at $12.1$ B	$\mathbf{-0.041}$	$+0.081$	$+0.046$	$\mathbf{-0.171}$	$\mathbf{-0.239}$	$\mathbf{-0.165}$
Equivalent extra training	$+19.6\%$	$+23.9\%^{\dagger}$	$+23.5\%^{\dagger}$	$+18.2\%$	$+8.3\%$	$+16.3\%$

Appendix

Table 25: At 345M, matching COAX requires substantially more baseline training. Results for the 345M-A147M multimodal model at seed 29 under a constant learning rate, evaluated without inference-time bias. The first row gives COAX’s relative validation-loss improvement at

10

B tokens. The second gives the extended baseline’s gap relative to the fixed COAX checkpoint after

12.1

B tokens; negative means the baseline has overtaken COAX. † denotes equivalence points extrapolated from the log-linear fit because I2T and T2I had not crossed within the measured training ladder.

Mixture-of-Experts Language Models

Mixture-of-Experts language models use input-dependent routing to activate a small subset of expert networks, increasing capacity without proportional compute.

See all 225 papers

UMoE:Unlocking Every Expert in Domain-Specific Training
Jul 13, 2026Xuefeng Li, Pengfei LiuMixture-of-Experts Pruning Mixture-of-Experts Language Models
1SII · 2SJTU · 3GAIR
Mixture of Experts for Low-Resource LLMs
May 17, 2026Ori Bar Joseph, Smadar Arvatz, Noam Kayzer +2Mixture-of-Experts Language Models Expert Routing
Do Domain-specific Experts exist in MoE-based LLMs?
Apr 7, 2026Giang Do, Hung Le, Truyen TranMixture-of-Experts Language Models Mixture-of-Experts Models
Applied Artificial Intelligence Intiative (A2I2), Deakin University

Jul 13, 2026J/K move · Enter open · S save

UMoE:Unlocking Every Expert in Domain-Specific Training

Xuefeng Li, Pengfei Liu

1SII · 2SJTU · 3GAIR

Mixture-of-Experts Pruning Mixture-of-Experts Language Models Domain Adaptation Mixture of Experts LLM Fine-Tuning Mixture-of-Experts Fine-Tuning+1

Open paper AskPDF arXiv

Mixture-of-Experts (MoE) models scale capacity without proportional compute cost and have become a key architecture for frontier large language models (LLMs). Yet domain-specific post-training inherits an expert pool shaped by mixed-domain pre-training: a substantial subset of experts contributes little on the target domain, and standard supervised fine-tuning (SFT) leaves the composition of this pool unchanged. We propose a simple, budget-preserving pipeline that realigns the expert pool to the target domain before fine-tuning. Given a target domain, we (1) prune the experts with lowest domain-aligned saliency, (2) regrow the expert pool to its original size through perturbation-based expert expansion, and (3) apply standard SFT. The resulting model preserves the original expert count, parameter count, and inference cost. With a single frozen recipe and no per-domain hyperparameter tuning, UMoE consistently improves over direct sft across two MoE architectures (Qwen3-30B-A3B and Qwen3.5-35B-A3B), five domains (math, code, science, tool-use, and agentic coding), and 12 benchmarks. Representative improvements are 3.4 points in math average accuracy, 6.0 points on SWE-bench Verified. On a strong in-house math corpus, direct sft already surpasses Qwen3-30B-A3B-Thinking (82.81 vs.\ 81.06), yet UMoE further raises the average to 84.17, an additional 1.36 points, demonstrating robustness to a substantially stronger SFT regime. Data-scaling experiments further show that the gain persists as training data grows. Analysis reveals that the direct-SFT model allocates substantial routed-expert compute to a low-saliency subset that can be removed post hoc with little average degradation; UMoE turns this redundant capacity into useful domain capacity and achieves lower training loss, with gains spanning all difficulty levels in downstream evaluation.

Abstract

Mixture-of-Experts Language Models

Explore similar work

UMoE:Unlocking Every Expert in Domain-Specific Training

Mixture of Experts for Low-Resource LLMs

Do Domain-specific Experts exist in MoE-based LLMs?