Source-linked AI summary
Overcoming Data Imbalances with Expert Selection Biases
Ayush K. Tarun, Clara Meister, Antoine Bosselut
TL;DR
Data imbalance can constrain the effective expert capacity available to minority categories, and prior category-aware specialization work had not established benefits for data-scarce categories. COAX applies category-conditioned router biases during pretraining to steer categories toward preferred expert subsets. It improves low-resource performance across modalities, languages, and code, with reported gains equivalent to additional training tokens at two model scales.
Problem
Data imbalance can concentrate minority-category tokens on too few experts or spread them inefficiently, and prior work had not established whether category-aware specialization benefits data-scarce categories.
Source: §1 · p. 1Method
COAX applies positive, category-conditioned biases to designated experts during pretraining, making each category’s preferred expert share an explicit allocation control while keeping all experts accessible to the router.
Source: §1 · pp. 1–2, §3.2 · p. 4Results
Across modalities, languages, and code, COAX improves low-resource performance; gains equal 18.5% and 2.7% additional training tokens at 345M-A147M and 10.6B-A1.2B, respectively.
Source: §1 · p. 2Takeaways & Limitations
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Data imbalance can leave minority categories with too little effective expert capacity, and earlier category-aware specialization work had not established benefits for data-scarce categories. COAX uses category-conditioned router biases to encourage category-aligned allocation and improves low-resource performance across modalities, languages, and code.
1 INTRODUCTION
- Data imbalance can concentrate minority-category tokens on too few experts or spread them inefficiently, limiting their available capacity.Prior multimodal and multilingual analyses report minority-modality collapse onto one expert and concentrated routing for underrepresented languages.Source: §1 · p. 1
- Earlier work encouraged category-aware specialization, but did not establish that it improves performance for data-scarce categories trained alongside higher-resource categories.The cited multimodal and multilingual studies report that low-resource performance can be worse than when training on the category alone.Source: §1 · p. 1
- COAX adds positive category-specific bias to designated experts during pretraining, letting a 1%-of-data language prefer 4 of 128 experts (3.1%) while retaining router override and standard inference.The preferred expert share can exceed a category’s share of training tokens, and the bias can be removed at inference.Source: §1 · p. 2, Figure 1 · p. 2
- Across modalities, languages, and code, COAX improves low-resource performance; gains equal 18.5% and 2.7% additional training tokens at 345M-A147M and 10.6B-A1.2B, respectively.The improvements often come with little to no degradation for higher-resource categories.Source: §1 · p. 2
2 RELATED WORK
Related work covers MoE routing and specialization, then distinguishes COAX by its explicit control over preferred expert allocation for categories. For low-resource MoEs, it contrasts COAX’s fixed active-expert count with methods that intervene through sharing, routing, or parameter allocation.
2 RELATED WORK
- Prior MoE work addresses expert under-utilization and routing collapse, while later analyses question whether experts develop human-interpretable expertise (Shazeer et al., 2017).Load-balancing biases have also been adjusted according to expert load (Wang et al., 2024; DeepSeek-AI, 2025).Source: §2 · pp. 2–3
- Category-aware methods use dedicated experts, language-guided routing, modality specialization, or routing-distribution objectives; COAX instead specifies a soft preferred expert subset and its size.Examples include DEMix (Gururangan et al., 2022), Lingual-SMoE, LIMoE (Mustafa et al., 2022), SMoES (Bo et al., 2026), and EDL (Li et al., 2026).Source: §2 · p. 3
- Low-resource approaches include increasing expert sharing, inference-time routing interventions, routing-alignment objectives, and allocating language-specific parameters after pretraining.LTDR increases activated experts for rare visual tokens, whereas COAX keeps active experts per token fixed and biases a larger expert subset during pretraining.Source: §2 · p. 3
3 METHOD
Sparse MoE routing selects a fixed number of experts per token, while COAX adds category-conditioned biases to steer selection toward designated expert subsets. Its preferred expert share can exceed a category’s token share without changing active experts per token; routing remains soft and standard at inference.
3.1 PRELIMINARIES: SPARSE MOE ROUTING
- Standard sparse MoE routing selects the K experts with highest affinities for each token and weights the selected experts by normalized affinity.The same routing operation is applied independently at each MoE layer; COAX leaves shared experts unchanged.Source: §3.1 · p. 3
3.2 COAX: CATEGORY-ORIENTED ALLOCATION OF EXPERTS
- COAX adds category-conditioned biases before top-K selection, steering tokens toward designated expert subsets while retaining access to other experts.The bias affects expert selection only; selected experts are weighted using the original routing scores, and the bias can be removed at inference.Source: §3.2 · pp. 3–4
- COAX is a soft routing prior, not a hard partition: tokens are encouraged toward their category’s subset, while cross-category expert sharing remains possible.An outside expert can still be selected if its learned score outweighs the bias, and experts in a category’s subset remain available to other categories.Source: §3.2 · p. 4
- A category’s preferred expert share can be set independently of its training-token share, allocating it a larger expert pool without increasing per-token active compute.For example, a category comprising 2.5% of training tokens can be biased toward 10% of routed experts, or a relative share of 4.Source: §3.2 · p. 4
4 EXPERIMENTAL SETUP
The experiments evaluate COAX in multimodal, multilingual, and multi-domain settings using three MoE model sizes. The setup specifies the model architecture, datasets, tokenizers, and common training configuration.
4 EXPERIMENTAL SETUP
- Experiments cover multimodal, multilingual, and multi-domain settings at 345M-A147M, 3B-A654M, and 10.6B-A1.2B model sizes.Training and evaluation use held-out splits from high-quality text, multilingual, vision, and audio datasets.Source: §4 · pp. 4–5
- All models use 32 layers, one shared expert, top-K = 8 routing, tied embeddings, SwiGLU, QK-Norm, and a sigmoid router.The models also use Grouped Query Attention.Source: §4 · p. 4
- Figure 2 measures the fraction of selected experts within a token’s designated subset; larger bias strengths consistently raise this in-group routing fraction.Dashed lines mark the expected fraction under uniform-random routing.Source: Figure 2 · p. 5
- Unless otherwise stated, models train on 10B tokens with sequence length 4096, auxiliary load-balancing loss, and COAX α = 0.03 for biased runs.Training uses 64 NVIDIA GH200 GPUs; expert subsets are mutually exclusive across categories.Source: §4 · p. 5
5 RESULTS
COAX steers category tokens toward designated expert subsets, with routing strength controlled by expert bias. Across multilingual, multimodal, and domain evaluations, it improves low-resource performance, while coarse category grouping and training-time bias can preserve gains without requiring bias at inference.
5 RESULTS
- Increasing α strengthens routing toward designated expert subsets, with in-group fractions peaking early, then decreasing before gradually rising again.The additive bias has comparatively greater influence early, when learned routing scores are close to uniform.Source: §5 · p. 5
5.1 MULTILINGUALITY
- Across the four-language mixture, gains grow as training-data share falls; Swahili improves substantially and significantly at every scale, while removing inference bias retains most benefits.Swahili is 0.85% of training data; Korean is 10% and receives relatively modest improvements. Table 1 reports mean relative validation-loss changes against unbiased baselines, with negative values indicating improvement.Source: §5.1 · pp. 5–6, Table 1 · p. 6
- COAX improves seven of eight languages in the 10.6B-A1.2B eight-language mixture, with Khmer approximately unchanged and English improving slightly.Relative loss changes include Tamil −0.42%, Georgian −0.58%, Urdu −0.31%, and Swahili −0.49%.Source: §5.1 · p. 6
5.2 MULTIMODALITY
- Vision shows the most consistent multimodal improvement across scales, with significant reductions in both image-to-text and text-to-image loss.At the smallest and largest scales, text also improves despite receiving a smaller fraction of the expert pool; most benefits persist without inference-time bias.Source: §5.2 · p. 7, Table 3 · p. 7
5.3 DOMAINS
- With inference bias retained, all four specialist domains improve at every scale, whereas without it only code consistently benefits across scales.The authors suggest non-code domains may share more useful structure with general English.Source: §5.3 · p. 8, Table 4 · p. 8
- In continual pretraining, COAX works best for code and next for vision, while newly introduced Russian and Hindi slightly underperform vanilla CPT; replay categories improve in all four settings.Source: §5.3 · p. 8
6 ABLATIONS
The ablations find that separating English from the long tail can work without fine-grained language partitions, and that bias strength trades minority-language gains against inference mismatch and overhead. Annealing bias to zero recovers much of the gain without inference-time bias.
6 ABLATIONS
- The English-versus-rest grouping performs as well as or better than per-language partitioning for six of eight languages.The comparison is for the 10.6B-A1.2B model on the eight-language mixture; the table reports relative validation-loss changes, with negative values better.Source: Table 5 · p. 8, §6 · p. 8
- For the 10.6B-A1.2B model, coarse grouping adds 1.3% training-time overhead per step, compared with 2.4% for finer grouping; the unbiased baseline needs 2.7% more training tokens to match its average validation loss.Source: §6 · p. 9
- With bias removed at evaluation, gains peak at α = 0.03 before mismatch becomes increasingly harmful; retaining bias yields monotonically greater minority-language gains, reaching mean ∆loss −1.65% at α = 0.3.Retaining stronger bias at inference increases throughput overhead; at α = 0.3 without inference bias, the mismatch is catastrophic.Source: §6 · p. 9, Table 6 · p. 9
- Annealing α = 0.3 to zero during cooldown recovers 76.4% of the minority-language improvement relative to keeping α = 0.3 active at inference.The minority-language changes are −1.26% versus −1.65%; English still degrades by +0.57%.Source: §6 · p. 9, Table 6 · p. 9
7 CONCLUSION
The conclusion presents expert biases as a controllable way to shape category-aligned routing and allocate preferred expert shares, while preserving soft sharing and standard inference. The experiments use sparse MoE models and high-quality training and evaluation datasets across several categories.
7 CONCLUSION
- Router biases introduced for load balancing can instead shape category-aligned routing and make preferred expert share an explicit training-time control.The conclusion says COAX need not use finely tuned expert partitions and can provide gains alongside conventional data upsampling.Source: §7 · p. 9
A.1 MODEL CONFIGURATIONS
- The experiments use highly sparse MoE configurations, with many routed experts but only a small subset activated per token.The three model scales are summarized in Table 7.Source: §A.1 · p. 15
A.2 DATASETS AND EVALUATION
- Training and evaluation use category-specific datasets, including English, multilingual, vision, audio, code, mathematics, legal, and biomedical text.Evaluation uses 100M held-out tokens per category except Swahili and interleaved audio-text, which use 20M each.Source: §A.2 · p. 15
B ADDITIONAL EVALUATIONS
- Additional evaluations cover task-level metrics beyond validation loss and continual pretraining in which COAX is introduced after initial pretraining.Source: B ADDITIONAL EVALUATIONS · p. 15
B.1 TASK-LEVEL EVALUATION
Teacher-forced task-level metrics improve with COAX, while multilingual downstream benchmarks are underpowered at this training scale. In continual pretraining, COAX improves vision, code, and replay on the original languages, but gains vary for new languages.
B.1 TASK-LEVEL EVALUATION
- Small per-language token budgets make standard multilingual downstream benchmarks unreliable for resolving the observed differences.MultiBLiMP differences are small and inconsistent across scoring schemes; SIB-200 accuracy is 25.10% for both models.Source: §B.1 · p. 15, Table 8 · p. 16
- COAX improves all twelve teacher-forced token- and word-level ASR and captioning metrics for the 345M-A147M model.The models are not sufficiently trained for meaningful free-running generation at this scale.Source: §B.1 · pp. 15–16
B.2 COAX IN CONTINUAL PRETRAINING
- Continual-pretraining results improve for vision, code, and replay on the original languages, but Russian and Hindi show essentially no new-category loss improvement.The experiment upcycles a 3B-A654M checkpoint from 128 to 160 routed experts and assigns the 32 new experts to a category comprising 20% of tokens.Source: §B.2 · p. 16
C ANALYSIS OF COAX-INDUCED SPECIALIZATION
The analyses test what COAX changes in routing and expert functionality, compare it with hard routing and divergence-based specialization, and examine design choices. They find functional localization without a sharper learned router, and favor partial specialization with sharing retained.
C ANALYSIS OF COAX-INDUCED SPECIALIZATION
- COAX improves all twelve teacher-forced token- and word-level ASR and captioning metrics in the 345M-A147M model.The metrics use reference transcripts or captions, rather than free-running generation.Source: Table 9 · p. 16
C.1 ROUTER ENTROPY UNDER COAX
- Router entropy ends at 4.492 for COAX versus 4.422 for the unbiased baseline, so gains do not arise from a sharper learned router.Entropy is measured before adding the category bias; COAX retains about 83% of the baseline’s entropy reduction.Source: §C.1 · p. 17, Figure 3 · p. 17
C.2 COAX INDUCES FUNCTIONALLY MODULAR EXPERTS
- Removing four Swahili-designated experts raises Swahili loss by 107.09% at α = 0.03 and 128.18% at annealed α = 0.3, versus 9.34% for the baseline.At α = 0.3, the same removal changes English, Korean, and Urdu loss by only +0.08%, +0.19%, and +0.03%, respectively.Source: §C.2 · pp. 17–18, Table 11 · p. 18
D.1 DEMIX-STYLE HARD ROUTING
- In the 10.6B eight-language setup, COAX reduced minority-language validation loss by 0.286% versus the same-seed unbiased baseline, with 1.79% iteration-time overhead versus hard routing’s 17.57%.The comparison is not perfectly matched: COAX’s bias is removed at evaluation, whereas the matched hard-routing condition retains its constraint.Source: §D.1 · p. 18, Table 12 · p. 19
D.2 EXPERT DIVERGENCE LEARNING
- At β = 5 × 10−4, EDL regresses the seven-language minority mean by +1.016%, while COAX improves it by −0.286%; greater router divergence alone does not improve the long tail.EDL directly optimizes router distributions, whereas COAX biases top-K selection and leaves probability-level divergence nearly unchanged.Source: §D.2 · p. 19, Table 13 · p. 20, Table 14 · p. 20
E.1 SENSITIVITY TO THE RELATIVE EXPERT SHARE
- For vision and audio, proportional expert capacity (rc = 1) consistently underperforms 2× preferred share (rc = 2), while tested Swahili allocations from 2 to 32 experts show little consistent trend.Additional capacity helps, but the tested range does not indicate a need to precisely tune the expert-set size.Source: §E.1 · p. 20, Table 16 · p. 21
E.2 BIASING ALL LAYERS IS A ROBUST DEFAULT
- Biasing all layers gives the best four-language mean (−0.158%), versus −0.029% for late layers, +0.105% for early layers, and +0.382% for the middle block.Language-selective routing increases with depth across the three model scales.Source: §E.2 · pp. 20–21, Table 17 · p. 22
E.3 INTERACTION WITH LOAD BALANCING
- Raising the load-balancing coefficient to 1e−2 reduces minority gains, while COAX changes the multimodal auxiliary loss by at most 0.31%.For α = 0.03, the minority-language gain shifts from −0.42% to −0.18%; for annealed α = 0.3, it shifts from −1.26% to −0.43%.Source: §E.3 · pp. 21–22, Table 18 · p. 22
E.4 COAX’S GAINS PERSIST WITH LONGER TRAINING
COAX’s advantage stabilizes rather than disappearing with longer training, including in a run extended to 100B tokens. The capacity-allocation tests find that biasing the majority category matters, while the single-seed comparison magnitudes are directional rather than definitive.
E.4 COAX’S GAINS PERSIST WITH LONGER TRAINING
- From 5B to 10B tokens, COAX reduces validation loss by 0.298 ± 0.028 percentage points on the six non-Khmer minority languages, with no downward trend.In a separate 100B-token run, the gain reaches 0.118 percentage points and remains stable after increasing slightly from 10B to 20B tokens.Source: §E.4 · pp. 22–23, Figure 4 · p. 23, Figure 5 · p. 23
- The 100B-token run uses lower-quality multilingual data, so its 10B-token gain is not directly comparable to Figure 4.The single-seed comparison treats magnitudes as directional rather than definitive.Source: §E.4 · p. 23, §F.1 · p. 24
F.1 BIASING THE MAJORITY CATEGORY IS IMPORTANT
- Removing the English bias reduces the mean minority improvement from −0.286% to −0.131%, losing to full COAX on seven of eight languages.The no-English-bias condition lets English use the entire expert pool for top-K selection.Source: §F.1 · p. 24, Table 20 · p. 24
F.2 REDUCING MAJORITY DATA DOES NOT REPRODUCE COAX
Reducing majority-category data does not reproduce COAX’s capacity reallocation, while data upsampling and COAX provide additional gains together. COAX’s throughput overhead decreases sharply with model scale.
F.2 REDUCING MAJORITY DATA DOES NOT REPRODUCE COAX
- COAX improves minority validation loss by −0.286%, whereas reducing English data degrades it by +2.015%; every language worsens in the data-reduction run.The comparison reduces total training tokens by 33.7%, so it rejects a simple equivalence but does not quantify relative efficiency.Source: §F.2 · pp. 24–25, Table 21 · p. 25
F.3 COAX AND DATA UPSAMPLING ARE COMPLEMENTARY
- With Urdu and Swahili doubled to 2 epochs, upsampling alone improves the minority mean by −3.579%, and combining it with COAX reaches −3.738%.The combination adds −0.159% over upsampling alone and captures approximately 61% of COAX’s standalone benefit; upsampling increases compute by 3.61%.Source: §F.3 · p. 25, Table 22 · p. 25
G.1 TRAINING THROUGHPUT OVERHEAD OF COAX
- COAX’s iteration-time overhead falls from approximately 20–22% at 345M-A147M to 1.7–2.4% at 10.6B-A1.2B.Throughput measurements average steps 8050–8200 from the step-8000 checkpoint; the authors hypothesize larger expert GEMMs better amortize routing and dispatch overheads.Source: G COMPUTATIONAL EFFICIENCY · p. 25, §G.1 · p. 26, Table 23 · p. 26
G.2 HOW MUCH ADDITIONAL TRAINING IS COAX WORTH?
The section translates COAX loss gains into equivalent additional baseline training, finding a larger equivalence at 345M than at 10.6B and a favorable measured training-cost trade-off at 10.6B with coarser expert grouping.
G.2 HOW MUCH ADDITIONAL TRAINING IS COAX WORTH?
- Equivalent additional training is measured by continuing the matching baseline at the same learning rate until it reaches the fixed COAX checkpoint's loss.The models are evaluated periodically on the same held-out sets; the first baseline match gives the equivalence.Source: §G.2 · p. 26
- The baseline required 18.5% additional training tokens to match weights-only COAX loss at 345M-A147M and 2.7% at 10.6B-A1.2B.The 345M result averages six validation sets; the 10.6B result uses the two-group eight-language configuration.Source: §G.2 · p. 26, Table 25 · p. 27, Table 24 · p. 27
- Keeping the COAX bias active at evaluation raises the equivalence to 20.8% at 345M and 3.8% at 10.6B, or 4.5% at 10.6B excluding Khmer.At 345M, the two vision sets had not crossed by 12.1B baseline tokens and required extrapolation individually.Source: §G.2 · pp. 26–27
- The equivalence falls from 18.5% at 345M to 2.7% at 10.6B because additional training reduces loss faster in the larger-model experiment.COAX throughput overhead falls from 20.2% to 2.4%, and from 2.4% to 1.3% with coarser expert grouping.Source: §G.2 · p. 27