Quantile Balancing
Summary: Quantile Balancing is an auxiliary-loss-free, hyperparameter-free MoE load balancing technique that derives expert allocation directly from router score quantiles — replacing heuristic auxiliary losses with a direct statistical constraint satisfied by construction.
Overview
Standard MoE balancing uses auxiliary losses (entropy, CV², switch loss) with a balancing hyperparameter λ. These require tuning and can conflict with task loss. Quantile Balancing eliminates this by making balanced allocation a constraint satisfied by construction rather than an objective optimized via gradients.
Key Components
Router Score Quantiles
- For each expert
e, collect router scores across a token batch:S_e = {s_{e,1}, ..., s_{e,B}} - Compute empirical quantiles:
Q_e(p) = quantile(S_e, p)for target probabilities - Target: uniform allocation → each expert should receive ~1/E of tokens
Allocation by Construction
- Instead of
top-k(softmax(router(x))), compute routing assignments that match the target quantile profile - Algorithm: sort tokens by router score per expert, assign to fill quantile bins
- Result: exact balance (up to batch granularity) without λ tuning
No Gradient Through Balancing
- Balancing is a discrete assignment step — no gradient flows through it
- Router gradients come purely from task loss
- Eliminates gradient conflict between task and balancing objectives
Mathematical Formulation
In plain English: Rather than penalizing imbalance after the fact, Quantile Balancing ensures each expert gets its fair share by construction — you rank tokens by their affinity for each expert and hand out assignments to perfectly fill each expert's quota.
Standard auxiliary loss:
# Router logits: [B, E]
router_logits = router(x)
probs = softmax(router_logits, dim=-1)
topk_probs, topk_idx = probs.topk(k, dim=-1)
# Auxiliary losses
entropy_loss = -(probs * log(probs)).sum(-1).mean()
cv_squared = (probs.mean(0) / probs.mean(0).mean()) ** 2 - 1
L_balance = λ * (entropy_loss + cv_squared)
Quantile Balancing:
# For each expert independently
for e in range(E):
scores_e = router_logits[:, e] # [B]
# Target: each expert gets B/E tokens (for top-1)
sorted_idx = scores_e.argsort(descending=True)
# Assign top B/E tokens to expert e
assignments[sorted_idx[:B//E]] = e
# No λ, no auxiliary loss, no gradient through assignment
# Router trained purely on task loss
Variants / Extensions
- Top-k Quantile Balancing: Extend to top-k routing with per-slot quantiles
- Hierarchical Quantiles: Balance at node level (expert parallel group) then within node
- Adaptive Quantiles: Non-uniform targets for heterogeneous expert capacities
Applications
- Stable LatentMoE (Kimi K3: 896 experts, 16 active)
- Any extreme-sparsity MoE (>256 experts, <2% activation)
- Expert parallel training on large clusters (64+ nodes)
- Scenarios where balancing hyperparameter tuning is infeasible
Historical Context
| Year | Work | Contribution |
|---|---|---|
| 2017 | MoE (Shazeer) | Auxiliary loss for balancing |
| 2021 | Switch Transformer | Simplified top-1 + aux loss |
| 2022 | GLaM / GShard | Large-scale, capacity factors |
| 2023 | DeepSeek-V2 | Fine-grained experts, aux-free attempts |
| 2024 | Expert Choice / Various | Reverse routing, capacity-free |
| 2026 | Quantile Balancing (Kimi K3) | Direct quantile constraint, zero hyperparameters, construction-based |
Related Concepts
- stable-latentmoe
- moe-routing
- expert-parallelism
- kimi-k3
- auxiliary-loss-free-moe
Sources
- kimi-k3-doc: Introduction — Quantile Balancing derives allocation from router-score quantiles, no heuristic updates/hyperparameters
- kimi-k3-tech-blog: Architecture and Infrastructure — Quantile Balancing eliminates heuristic updates and sensitive balancing hyperparameter