# Overcoming Data Imbalances with Expert Selection Biases

> COAX repurposes router biases as category-conditioned routing priors, boosting low-resource performance across languages, modalities, and domains without inference-time overhead.

- **Source:** [arXiv](https://arxiv.org/abs/2610.13370)
- **Published:** 2026-10-31
- **Permalink:** https://picx.dev/p/zMcEwE
- **Whiteboard:** https://picx.dev/p/zMcEwE/image

## Summary

# Overcoming Data Imbalances with Expert Selection Biases

## Summary (Overview)

- **Core Contribution**: The paper introduces **COAX (Category-Oriented Allocation of eXperts)**, a lightweight method that repurposes per-expert router biases (originally designed for load balancing) as category-conditioned routing priors during pretraining, enabling explicit control over expert allocation for low-resource categories.

- **Mechanism**: COAX adds a positive bias $\alpha$ to router scores of a designated expert subset $P_c$ for each category $c$ during pretraining only, softly encouraging category-aligned specialization while preserving cross-category knowledge sharing and requiring no inference-time modifications.

- **Key Concept — Preferred Expert Share**: COAX allows setting a category's preferred fraction of the expert pool ($q_c = M_c/N$) independently of its training data share ($p_c$), enabling over-allocation of expert capacity to data-scarce categories via the relative expert share $r_c = q_c/p_c$.

- **Results**: Across modalities, languages, and domains at scales up to 10.6B parameters, COAX consistently improves low-resource performance with minimal impact on high-resource categories. Gains are equivalent to **18.5%** and **2.7%** additional baseline training tokens at the 345M and 10.6B scales, respectively.

- **Practical Advantages**: COAX requires no changes to the standard MoE inference stack, no category labels at inference, and its training overhead decreases with scale (≤1.3% at 10.6B parameters).

## Introduction and Theoretical Foundation

### Background

Mixture-of-Experts (MoE) models have become the standard architecture for frontier LLMs, with multiple parallel experts in each transformer block theoretically enabling expert specialization according to domains, languages, or modalities. However, two key problems undermine this potential:

1. **Category-aligned specialization does not reliably emerge** under standard training (Jiang et al., 2024; Wang et al., 2026).
2. **Data imbalance causes asymmetric capacity allocation**: tokens from minority modalities can collapse onto a single expert and saturate its capacity (Mustafa et al., 2022), and underrepresented languages exhibit concentrated routing (Joseph et al., 2026).

### Motivation

For low-resource categories, training alongside higher-resource categories degrades performance relative to training on the category alone, as documented in multimodal (Aghajanyan et al., 2023; Huang et al., 2022) and multilingual (Longpre et al., 2025) settings. Prior work on category-aware specialization (DEMix, Lingual-SMoE, EDL) established that category information can shape routing but did not demonstrate that specialization can aid data-scarce categories.

### Theoretical Foundation

COAX builds on the observation that expert-wise routing biases (Wang et al., 2024) can be repurposed from load balancing to category-conditioned routing priors. The key insight is that a **soft bias**—rather than a hard partition—can encourage specialization while retaining all experts' accessibility, preserving the possibility of cross-category knowledge sharing.

## Methodology

### Preliminaries: Sparse MoE Routing

For a token $t$ with input representation $u_t \in \mathbb{R}^d$, the router computes affinity with expert $i$ as:

$$s_{i,t} = G\left(u_t^\top e_i\right), \quad i \in \{1, \ldots, N\}$$

where $e_i$ is the learned router representation of expert $i$ and $G(\cdot)$ is the nonlinear gating function. Standard top-$K$ routing selects:

$$A_t = \text{TopK}(s_t, K)$$

with normalized expert weights:

$$g_{i,t} = \frac{s_{i,t} \mathbb{1}\{i \in A_t\}}{\sum_{j \in A_t} s_{j,t}}$$

The MoE block output is:

$$h_t = u_t + \sum_{i=1}^{N} g_{i,t} \text{FFN}_i(u_t)$$

### COAX: Category-Oriented Allocation of Experts

For each category $c$, COAX designates a subset of routed experts:

$$P_c \subseteq \{1, \ldots, N\}, \quad M_c = |P_c|$$

For a fixed bias $\alpha \geq 0$ shared across categories:

$$b_{c,i} = \alpha \mathbb{1}\{i \in P_c\}$$

During pretraining, selection scores become:

$$\tilde{s}_{i,t} = s_{i,t} + b_{c_t, i}$$

yielding selected experts $\tilde{A}_t = \text{TopK}(\tilde{s}_t, K)$. Importantly, **the bias affects only expert selection**—once top-$K$ experts are selected, their contributions use the original unbiased scores:

$$g^{\text{COAX}}_{i,t} = \frac{s_{i,t} \mathbb{1}\{i \in \tilde{A}_t\}}{\sum_{j \in \tilde{A}_t} s_{j,t}}$$

### Key Definitions

**Preferred expert share** (fraction of routed experts toward which category $c$ is biased):

$$q_c = \frac{M_c}{N}$$

**Relative expert share** (preferred share divided by token share $p_c$):

$$r_c = \frac{q_c}{p_c} = \frac{M_c}{N p_c}$$

When $r_c > 1$, the category is biased toward a larger fraction of experts than its data prevalence. Increasing $M_c$ does **not** change the number $K$ of active experts per token, so per-token compute stays constant.

### Experimental Setup

- **Model scales**: 345M-A147M, 3B-A654M, and 10.6B-A1.2B (all with 32 layers, GQA, one shared expert, top-$K=8$)
- **Training**: 10B tokens (unless otherwise stated), 64 NVIDIA GH200 GPUs, Megatron-LM framework, WSD scheduler
- **Default hyperparameters**: $\alpha = 0.03$, mutually exclusive expert subsets across categories
- **Settings**: Multilingual (English/Korean/Urdu/Swahili and 8-language mixtures), multimodal (text/vision/audio), and multi-domain (English/code/math/legal/medical)

## Empirical Validation / Results

### Multilinguality

On a four-language mixture (86.42% English, 10% Korean, 2.73% Urdu, 0.85% Swahili) with relative expert shares of 2, 4, and 4 for Korean, Urdu, and Swahili:

| Model | English | Korean | Urdu | Swahili |
|-------|---------|--------|------|---------|
| 345M-A147M | +0.05 | −0.01 | −0.23 | −0.55** |
| 3B-A654M | −0.01 | −0.11 | −0.31* | −0.40** |
| 10.6B-A1.2B | +0.02 | −0.00 | −0.10* | −0.30** |

*Table: Mean relative change in validation loss (∆%) vs. unbiased baselines (negative is better). **Bold** indicates improvement; * and ** denote 90% and 95% confidence intervals excluding zero.*

**Key findings**:
- Gains **increase as languages become lower-resource** (Swahili at 0.85% benefits most)
- Removing the bias at inference preserves most of the benefit
- COAX induces **functionally localized experts**: disabling the four Swahili-designated experts increases Swahili loss by 107.1% vs. only 9.3% for the baseline's most-used experts

### Comparison with Alternative Methods (8-language mixture)

| Method | Mean minority Δ loss |
|--------|---------------------|
| **COAX** | **−0.286%** |
| Hard-routing (DEMix-style) | +0.047% |
| EDL | +1.016% |

COAX achieves the best results, highlighting the importance of **higher relative expert shares** combined with soft specialization.

### Multimodality

On a mixture of 75% text, 12.5% vision, and 12.5% audio with $r_c = 2$ for vision and audio:

- **Vision shows the most consistent improvements** (both I2T and T2I) across all scales
- At 345M-A147M and 10.6B-A1.2B, text also improves despite receiving a smaller expert share
- Gains persist without inference-time bias, indicating the training intervention induces useful specialization

### Domains

With English under-provisioned ($r_c \approx 0.7$), code and math at $r_c = 2$, and legal/medical at $r_c = 4$:
- Without inference bias, only code consistently benefits
- **With inference bias retained, all four specialist domains improve at every scale**

### Ablations

**Coarse expert grouping suffices**: A simple two-group split (English vs. all non-English) performs as well as or better than per-language partitioning for 6 of 8 languages, suggesting much of the benefit comes from separating the dominant language from the long tail.

**Bias strength ablation** (3B-A654M model):

| α | English (no bias) | Minority mean (no bias) | Minority mean (with bias) |
|---|---|---|---|
| 0.003 | −0.02 | −0.12 | −0.12 |
| 0.01 | −0.10 | −0.30 | −0.31 |
| 0.03 | −0.03 | −0.42 | −0.48 |
| 0.1 | +0.38 | −0.37 | −0.96 |
| 0.3 | +1.88 | +19.32 | −1.65 |

Without inference bias, gains peak at $\alpha = 0.03$ before train-test mismatch dominates. **Annealing $\alpha$ from 0.3 to zero during cooldown recovers 76.4% of the minority-language gain** without requiring inference-time biasing.

## Theoretical and Practical Implications

### Theoretical Implications

1. **Router biases serve a broader role than load balancing**: The results demonstrate that expert biases can shape category-level allocation, providing a lightweight control mechanism for expert specialization.

2. **Soft specialization outperforms hard partitioning**: COAX's soft bias outperforms DEMix-style hard routing and EDL, suggesting that retaining cross-category sharing while encouraging specialization is beneficial.

3. **Interference reduction, not isolation, drives gains**: The coarse grouping results suggest that much of COAX's benefit comes from reducing interference between dominant and minority categories rather than isolating every minority category.

4. **Functional localization emerges**: COAX induces experts that are functionally specialized (disabling them is highly damaging for their assigned category), unlike standard MoE training where routing reflects representation geometry rather than semantic categories.

### Practical Implications

1. **Zero inference overhead**: The bias can be removed at inference, yielding a standard top-$K$ MoE with no category-label requirements.

2. **Scalable training cost**: Training overhead decreases with scale (≤1.3% at 10.6B-A1.2B), and the performance gains translate to 2.7% additional training tokens at this scale.

3. **Complementary to data upsampling**: COAX provides additional gains when combined with conventional data upsampling strategies.

4. **No fine-tuning required**: Coarse expert groupings suffice, and performance is relatively insensitive to precise expert share once $r_c \geq 2$.

## Conclusion

COAX demonstrates that expert biases can serve as a lightweight control for category-level expert allocation during pretraining, making room for low-resource categories while retaining soft cross-category sharing. The method:

- Improves low-resource performance across modalities, languages, and domains
- Preserves standard MoE inference with no additional overhead
- Requires no finely tuned expert partitions
- Provides favorable trade-offs between performance gains and training overhead at scale

**Future directions** include replacing manually specified categories with learned structure and exploring systems-aware biasing that aligns preferred experts with expert-parallel topology for improved training efficiency.

---

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