N Load‑Balancing Techniques for Large‑Scale MoE Models

MoE load imbalance

Mixture‑of‑Experts (MoE) layers route each token to a subset of experts. For a batch with T tokens and E routed experts, the router produces logits, applies softmax or sigmoid to obtain routing scores, and then selects the top‑k experts. The actual token load of an expert is the sum of the selected tokens' scores. Without explicit balancing a small random advantage early in training can be amplified, causing a few experts to become overloaded while others starve – a phenomenon known as routing collapse.

Consequences of load imbalance:

Computation bottleneck: in synchronous training the overall speed is limited by the busiest device.

Token overflow: experts have limited capacity; excess tokens must be dropped or re‑routed.

Expert degradation: under‑utilized experts receive insufficient training, wasting model capacity.

1. Importance loss

The 2017 Sparsely‑Gated MoE paper defines an expert’s Importance as the sum of its gating weights across the batch. Example with three experts and four tokens:

Token 1 → Expert 1: 0.9
Token 2 → Expert 1: 0.9
Token 3 → Expert 2: 0.2
Token 4 → Expert 1: 0.9

The resulting importance vector is [1.8, 0.2, 0], showing a clear bias toward Expert 1. To penalize imbalance, the squared coefficient of variation (CV) of the importance vector is added to the loss. CV = 0 when all importances are equal and grows as the distribution becomes more skewed. Importance measures only total gating weight, not the actual number of tokens. Two experts can have identical importance but very different token counts (e.g., Expert A receives 2 tokens with weight 0.9 each, Expert B receives 9 tokens with weight 0.2 each). Therefore, importance loss alone is insufficient.

2. Expected load loss

Original Sparsely‑Gated MoE uses a Noisy Top‑k gate, making the exact count of selected tokens nondifferentiable. By adding Gaussian noise to the logits, the probability that an expert appears in the top‑k can be estimated, yielding a differentiable expected load. Consider three experts with a batch of four tokens and Top‑1 routing. After adding noise, the per‑token probabilities might be:

Token 1: Expert 1 0.80, Expert 2 0.15, Expert 3 0.05

Token 2: Expert 1 0.70, Expert 2 0.20, Expert 3 0.10

Token 3: Expert 1 0.10, Expert 2 0.60, Expert 3 0.30

Token 4: Expert 1 0.75, Expert 2 0.15, Expert 3 0.10

Summing these probabilities gives the expected load vector [2.35, 1.10, 0.55]. Expected Load focuses on the *number* of tokens an expert is likely to receive, complementing Importance which focuses on weight magnitude. Early MoE models often combine both losses.

3. f·P auxiliary loss

Switch Transformer and GShard introduce a compact auxiliary loss that jointly considers:

f : the hard count of tokens actually dispatched to an expert (non‑differentiable).

P : the average soft probability the router assigns to that expert (continuous).

The loss penalizes large discrepancies between f and P. Keeping both terms is useful because f reflects the true discrete routing result while P provides gradients to the router. If an expert’s P becomes too large, the loss pushes the router to lower its selection probability. In the ideal state, f and P are close.

The hyper‑parameter controlling the loss strength must be balanced: too small leaves the router biased, too large interferes with the primary language‑model objective.

4. Multi‑level auxiliary loss

Modern models often apply three hierarchical balances:

Expert‑level balance : prevents individual experts from overheating or starving.

Device‑level balance : aggregates the load of all experts on a GPU and encourages each device to have similar total compute.

Communication‑level balance : limits the amount of token data sent across devices during the All‑to‑All exchange, avoiding network hotspots.

DeepSeek‑V2 explicitly discusses all three; whether to enable each depends on expert placement, parallel scale, and network topology.

5. Expert capacity + token dropping

Auxiliary losses only “encourage” balance and cannot guarantee that every token is processed evenly. Production systems therefore assign each expert a capacity factor (CF) . When the router requests more tokens than an expert’s capacity, common strategies are:

Skip the expert and follow the residual branch.

Drop the excess tokens.

Re‑route them to the next‑best expert.

Expand the shape dynamically (at the cost of predictable throughput).

A small CF leaves little buffer, increasing token drops; a large CF reduces drops but adds padding, memory usage, and worst‑case compute.

6. Expert choice

Traditional token‑centric routing selects the top‑k experts for each token, guaranteeing each token activates k experts but not that each expert receives a balanced token count. Expert Choice inverts the problem: each expert independently selects the top‑C tokens with the highest scores from a shared candidate pool. This fixes a bucket per expert, ensuring a deterministic token count per expert.

Side effects:

Expert compute becomes fixed, while token compute varies (some tokens may be selected by multiple experts, others by none).

The selection requires a global view of the token set, which is fine for bidirectional encoders or vision tasks but can cause future‑token leakage in causal language models.

7. Balanced assignment

Balanced Assignment formulates routing as a linear assignment problem: given an affinity matrix A_{ij} (score of assigning token i to expert j ), maximize total affinity while enforcing that each token is assigned to exactly one expert and each expert receives a predetermined number of tokens. When exact discrete solving is too costly, a Sinkhorn‑style row‑column normalization produces a soft assignment that can be discretized afterwards.

Pros: clear balance objective and strict constraints. Cons: requires global coordination across a token batch, adds extra computation and communication, and must respect causality in autoregressive models.

8. Auxiliary‑loss‑free (loss‑free bias)

Instead of back‑propagating a balance loss, this approach injects a per‑expert bias before the Top‑k selection. After each training step the actual load of the previous batch is measured; the bias for overloaded experts is decreased, while the bias for under‑loaded experts is increased. The bias only influences the discrete Top‑k decision and is **not** added to the final expert output weights.

Overloaded expert → bias ↓

Under‑loaded expert → bias ↑

Balanced load → bias unchanged

This online controller avoids extra gradient interference with the language‑model loss. Updates must be based on historical load, not on future tokens, to preserve causality. DeepSeek‑V3 adopts this loss‑free bias together with a tiny sequence‑wise balance loss to prevent extreme per‑sequence skew.

Summary of methods

Soft regularizers (Importance, Expected Load, f·P, Multi‑level) do not guarantee exact balance and require tuning of loss coefficients.

Hard caps (Capacity + Token Dropping) enforce an upper bound but may still leave some experts idle.

Routing changes (Expert Choice, Balanced Assignment) can enforce exact or near‑exact token counts, at the cost of variable per‑token compute or additional global coordination.

Online control (Loss‑Free Bias) offers near‑balance without extra gradients, but introduces a feedback latency that must be tuned.

Choosing a strategy depends on the model’s scale, hardware topology, and tolerance for extra computation or communication overhead.