Low-resource languages stop losing the expert pool: a category bias on routing, removable at inference

In a mixture-of-experts model, each category can be assigned a designated subset of experts, and training steers that category's tokens toward it, so a low-resource category gets more expert capacity than its share of the data. Earlier ways of making experts divide by category either gave each domain dedicated experts or used an objective to push category-level routing distributions apart; those results did not establish that category separation actually helps data-scarce categories. On a mixture of 86.42% English, 10% Korean, 2.73% Urdu, and 0.85% Swahili, gains grow as a language's data share falls: Korean improves modestly, Urdu more consistently, and Swahili substantially at every model scale.

The mechanism reuses the per-expert routing bias originally introduced for load balancing, but changes what the bias depends on from whether an expert is underused to which category a token belongs to. A load-balancing bias is adjusted dynamically by expert utilization to even out expert usage; COAX instead adds one fixed positive bias to a designated expert subset per category, turning the preferred expert share into an explicit training-time control. The bias affects only top-k selection, and selected experts are still weighted by the original unbiased routing scores, so it acts as a soft prior rather than a hard partition; it was tested at 345M-A147M, 3B-A654M, and 10.6B-A1.2B on multilingual and multimodal mixtures.

Removing the category bias at inference returns the model to standard top-k MoE, with no category labels and no change to the inference stack. Training with the bias and inferring without it creates a train-test mismatch, and that mismatch becomes harmful once the bias strength is large. Dropping the inference-time bias still preserves most of the benefit; annealing the bias linearly to zero during cooldown recovers most of the minority-language improvement, though English still degrades, so the dominant-language trade-off is not eliminated.

The method was validated at three scales, 345M-A147M, 3B-A654M, and 10.6B-A1.2B, on multilingual, multimodal, and multi-domain mixtures trained on 10B tokens, with held-out loss as the main metric and top-k accuracy on ASR and captioning as supporting evidence for the multimodal part. It targets model developers controlling expert capacity allocation during pretraining, and inference remains standard top-k MoE without category labels. The source proposes replacing manually specified categories with a learned structure and exploring systems-aware biasing aligned with the expert-parallel topology.

As bias strength grows, the mismatch between biased training and unbiased inference becomes increasingly harmful, and annealing to zero recovers most but not all minority-language improvement, with English still degrading. Multilingual evaluation relies mainly on held-out loss because standard downstream benchmarks are underpowered at this training scale: MultiBLiMP differences are small and inconsistent across scoring variants, Belebele sits near chance, and SIB-200 gives identical accuracy for both models. Free-running generation remains undertrained, with greedy ASR word error rate high and captioning CIDEr-D near zero, so task-level conclusions hold only under teacher forcing. In continual pretraining, COAX helps for code and vision but newly introduced Russian and Hindi slightly underperform vanilla continual pretraining.