Mixture-of-Experts (MoE) Routing and Load Balancing: Mathematical Foundations of Noisy Top-k Gating, Expert Capacity, Auxiliary Balancing Losses, and Loss-Free Load Balancing
Scaling dense transformer models incurs computational costs that grow proportionally with parameter count. In dense architectures, every token in a sequence activates all feed-forward network (FFN) parameters across every layer. Sparse Mixture-of-Experts (MoE) architectures decouple parameter count from per-token compute by partitioning dense FFN layers into multiple independent sub-networks ("experts") and deploying a learned gating router to direct each token to a sparse subset of experts.
While sparse activation allows models to scale to hundreds of billions or trillions of parameters at the inference and training cost of a much smaller dense model, it introduces a severe optimization instability: routing collapse. Left unconstrained, learned routers naturally degenerate into trivial attractor states where a small fraction of experts receive nearly all tokens, while the remaining experts receive zero gradients and starve.
Addressing routing collapse has driven a ten-year progression in neural architecture design, spanning stochastic noise injection, capacity-constrained communication buffers, auxiliary load-balancing objectives, router stability losses, and recently, auxiliary-loss-free dynamic bias adjustments.

1. Mathematical Formulation of Sparse MoE Gating
In a standard transformer layer, the token representation passes through a Multi-Head Attention (MHA) block followed by a Feed-Forward Network (FFN):
In a sparse MoE layer, the monolithic FFN is replaced by a collection of parallel expert networks , where each expert possesses independent parameter weights . A parametric routing module assigns a sparse weighting vector over all experts.
The layer output is computed as the linearly weighted combination of the selected expert outputs:
Because is constrained to be sparse (having at most non-zero entries), only expert forward passes are executed per token.
Classical Softmax Routing and Its Failure Mode
The naive implementation of a router computes affine logits followed by a standard softmax distribution:
where is the learnable gating weight matrix. To enforce sparsity, one selects the indices corresponding to the top- entries in .
Under naive gradient descent, this mechanism exhibits positive feedback loops:
- If expert randomly initializes with slightly better representation alignment for a cluster of tokens, it receives higher gating weights .
- Because expert processes more tokens, its weights receive frequent gradient updates, improving its capability faster than dormant experts.
- The router observes that expert minimizes loss more effectively, further increasing gating probabilities toward expert .
- Dormant experts receive fewer tokens, fewer gradients, and eventually permanently shut down.
2. Noisy Top-k Gating
To break the deterministic positive feedback loop during training, Shazeer et al. (2017) introduced Noisy Top-k Gating. The router injects tunable Gaussian noise into the unnormalized gating logits prior to top- selection.
Given input , the router evaluates a base affinity logit and a learned per-expert noise scale:
where represents the primary gating projection, controls the input-dependent noise magnitude, and guarantees non-negative standard deviations.
The router then applies a thresholding operation that preserves the largest values and sets all remaining coordinates to :
The final gating distribution is obtained via softmax normalization over the thresholded logits:
Because elements outside the top- are set to , their softmax probabilities evaluate to exactly zero:
where .
Gradient Flow Through Noisy Gating
The non-zero gating weights for are smooth and differentiable with respect to , , and . The gradient of the task loss with respect to the gating logit for an active expert is:
While the discrete set selection introduces non-differentiable boundary transitions, the stochastic noise term ensures continuous exploration of alternative expert subsets during early optimization.
3. Distributed Hardware Constraints: Expert Capacity and Buffer Overflow
In large-scale distributed training, experts are distributed across separate GPU devices using Expert Parallelism (EP). To maximize hardware utilization on tensor accelerators (NVIDIA GPUs, Google TPUs), computation kernels require static, predetermined tensor dimensions.
If tokens were dynamically routed without limits, an overloaded expert would exceed its allocated memory buffer, causing Out-Of-Memory (OOM) faults or stalling communication collectives.
To guarantee static tensor shapes, Lepikhin et al. (GShard, 2020) and Fedus et al. (Switch Transformers, 2022) formalized the concept of Expert Capacity.
Mathematical Definition of Expert Capacity
Given a batch of sequences, each with length , the total number of tokens entering an MoE layer is . If each token is routed to experts, the total number of token-expert assignments is .
Under uniform distribution across experts, each expert expects to receive exactly tokens. The Expert Capacity is defined by scaling this baseline by a capacity factor :
Batch of T Tokens (T = B x S)
│
▼
Learned Gating Router: Top-k(H(x))
│
┌──────────────────────────┼──────────────────────────┐
▼ ▼ ▼
Expert 1 Buffer Expert 2 Buffer Expert N Buffer
Capacity: C slots Capacity: C slots Capacity: C slots
┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐
│ Token 1 │ │ Token 3 │ │ Token 2 │
│ Token 4 │ │ Token 7 │ │ Token 5 │
│ [Pad Slot] │ │ Token 9 │ │ Token 8 │
└─────────────────┘ └─────────────────┘ └─────────────────┘
│
(If Tokens > C: OVERFLOW)
▼
Dropped Token Bypass (Residual: y = x)Buffer Overflow and Token Dropping
Each expert allocates a static input buffer of shape . When tokens are routed to expert :
- If the number of tokens assigned to expert is less than or equal to , the remaining slots are filled with zero-padding tokens. Padding tokens are ignored during loss computation and gradient backpropagation.
- If the number of assigned tokens exceeds , the excess tokens are dropped. A dropped token bypasses the expert computation entirely and is passed directly to the next layer via the identity residual connection ().
The capacity factor embodies a fundamental engineering trade-off:
- Low Capacity Factor (): Zero memory waste on padding, but high token drop rates whenever routing distribution drifts from uniform balance. Dropped tokens degrade language modeling perplexity.
- High Capacity Factor (): Eliminates token dropping at the expense of allocating substantial GPU memory to idle zero-padding, reducing effective batch throughput.
4. Auxiliary Load Balancing Losses
Because token dropping hurts model convergence, MoE architectures introduce auxiliary regularization terms to force the gating network toward uniform token distribution across all available experts.
The Switch Transformer Auxiliary Loss
In the Switch Transformer, Fedus et al. formulated a clean, differentiable auxiliary loss. For a given batch containing tokens, two metrics are computed for each expert :
- Fraction of Routed Tokens (): The proportion of total routing decisions assigned to expert :
where is the indicator function.
- Average Routing Probability (): The mean softmax gating probability allocated to expert across all tokens in the batch:
The auxiliary load balancing loss is defined as the scaled inner product of the vectors and :
where is a hyperparameter (typically ).
Proof of Uniform Minimum
The product couples the non-differentiable assignment count with the differentiable probability mass .
By the Cauchy-Schwarz inequality, for any non-negative probability vectors (where and ):
When , this simplifies to:
The global minimum of is achieved if and only if:
At this balanced point:
If any expert attracts disproportionate traffic (e.g., ), both and increase simultaneously, driving up quadratically and penalizing the router weights .
Imbalance State: Expert 1 gets 70% of tokens, Experts 2-4 get 10% each
f = [0.70, 0.10, 0.10, 0.10], P = [0.70, 0.10, 0.10, 0.10]
L_aux = α * 4 * (0.49 + 0.01 + 0.01 + 0.01) = α * 4 * 0.52 = 2.08 α
Balanced State: Experts 1-4 get 25% of tokens each
f = [0.25, 0.25, 0.25, 0.25], P = [0.25, 0.25, 0.25, 0.25]
L_aux = α * 4 * (0.0625 * 4) = α * 4 * 0.25 = 1.00 α5. Router Numerical Stability: The Z-Loss
During large-scale pre-training of deep MoE models (such as ST-MoE, Zoph et al., 2022), routers frequently experience logit drift. Because the softmax operation is translation invariant:
the routing logits can grow to extreme positive values () without changing the resulting softmax probabilities. In 16-bit floating point formats (FP16 or BF16), large logit magnitudes cause catastrophic exponent overflow, yielding NaN activations that corrupt training checkpoints.
To maintain logit magnitudes within stable numerical ranges, ST-MoE introduced the Router Z-Loss:
where is a small weighting coefficient (typically ).
The Z-loss directly penalizes the logarithm of the softmax partition function (the log-sum-exp of logits). By driving large logit values toward zero, the Z-loss eliminates floating-point round-off errors and prevents gradient divergence across billions of tokens.
6. Token-Choice vs. Expert-Choice Routing
All top- architectures discussed above operate under Token-Choice Routing: individual tokens evaluate all experts and select the top- candidates.
In 2022, Zhou et al. (Google Brain) inverted this formulation by introducing Expert-Choice Routing.
Token-Choice Routing:
Each token selects Top-k experts.
Result: Variable expert load (some experts overflow, others starve).
Remedy: Capacity factor padding and auxiliary balancing losses.
Expert-Choice Routing:
Each expert selects Top-C tokens.
Result: Perfect expert load balance by construction (exactly C tokens per expert).
Consequence: Variable token allocation (some tokens get multiple experts, some get none).Expert-Choice Mathematical Formulation
Let represent the concatenated matrix of all tokens in a batch. The router computes a token-to-expert affinity score matrix :
where the softmax is computed column-wise across the token dimension , rather than row-wise across experts:
Each expert selects the top- tokens with the highest affinity scores in column :
Because every expert selects exactly tokens, every expert buffer is filled to 100% capacity with zero padding and zero dropped tokens.
Trade-Offs of Expert-Choice Routing
While Expert-Choice routing achieves optimal FLOP efficiency and hardware utilization, it introduces two fundamental architectural constraints:
- Uneven Token Representation: Important tokens (e.g., subject nouns or syntax roots) may be selected by 5 or 6 experts simultaneously, while simple tokens (e.g., punctuation or articles) may be selected by 0 experts, relying solely on residual connections.
- Causal Autoregressive Decoding Incompatibility: During autoregressive inference (step-by-step generation), future tokens do not exist. An expert cannot compute the top- tokens across a batch until all tokens have arrived. Consequently, Expert-Choice routing is well-suited for encoder and prefill models, but cannot be applied directly to standard causal next-token generation loops.
7. Fine-Grained Segmentation and Shared Experts
To improve expert specialization without inflating routing complexity, modern architectures have transitioned from coarse experts to fine-grained expert segmentation combined with dedicated shared experts. This paradigm was formalized in DeepSeekMoE (Dai et al., 2024).
Fine-Grained Expert Partitioning
Instead of deploying large experts with hidden dimension and routing top-, DeepSeekMoE splits each expert into smaller sub-experts with intermediate dimension . The total number of experts becomes , and the router selects top- sub-experts:
Standard MoE (e.g., Mixtral 8x7B):
Total Experts: N = 8
Active Experts: k = 2
Combinations: C(8, 2) = 28 routing paths
Fine-Grained MoE (DeepSeekMoE, m = 4):
Total Experts: 4 * 8 = 32
Active Experts: 4 * 2 = 8
Combinations: C(32, 8) = 10,518,300 routing pathsWith identical parameter count and computational FLOPs per forward pass, fine-grained partitioning expands the combinatorial capacity of the model by several orders of magnitude, allowing sub-experts to learn highly focused, disentangled features.
Dedicated Shared Experts
In standard MoE models, redundant common knowledge (e.g., English grammar rules, punctuation syntax, common code structure) must be duplicated across all independent experts to prevent performance degradation when routed.
DeepSeekMoE isolates this invariant knowledge by reserving experts that are always activated for every token, irrespective of router scores. The remaining slots are dynamically allocated via top- gating:
By delegating general-purpose representations to fixed shared pathways, dynamic routed experts specialize purely in niche domain capabilities (e.g., specific programming languages, mathematical theorems, or multilingual translations).
8. Auxiliary-Loss-Free Load Balancing
While auxiliary load balancing losses () prevent expert starvation, they introduce a fundamental structural conflict:
The gradient of the language modeling loss attempts to route tokens to the expert most capable of predicting the next token. Simultaneously, the gradient of the auxiliary loss forces the router to route tokens away from capable experts to equalize traffic.
When is too small (), load balancing fails, causing expert collapse and token dropping. When is too large (), gradient interference severely degrades validation perplexity.
To resolve this conflict, Wang et al. (DeepSeek, 2024) developed Auxiliary-Loss-Free Load Balancing, deployed in DeepSeek-V3.
┌──────────────────────────────────────┐
│ Input Token Representation x ∈ R^d │
└──────────────────┬───────────────────┘
│
┌─────────────┴─────────────┐
│ Affine Projection: x * Wg │
└─────────────┬─────────────┘
│
┌─────────────────────┴─────────────────────┐
│ │
▼ ▼
Logit + Dynamic Bias b_i Raw Affinity Logits (Unbiased)
(Used ONLY for Expert Selection) (Used ONLY for Softmax Gating Weights)
│ │
▼ │
Select Top-k Indices K │
K = TopK(x * Wg + b, k) │
│ │
└─────────────────────┬─────────────────────┘
│
▼
Compute Normalized Gating Probabilities
G(x)_i = Softmax_{j ∈ K}( (x * Wg)_j )
│
▼
Zero Auxiliary Loss Gradients on RouterDynamic Bias Adjustment Mechanism
Instead of adding a penalty term to the backward loss graph, the router introduces a learnable per-expert bias vector that is updated dynamically based on real-time hardware queue statistics.
- Top-k Selection with Bias: When routing token , the top- expert indices are selected using the bias-augmented affinity scores:
- Gating Weight Calculation without Bias: Once the expert indices are selected, the actual gating weights used to multiply expert outputs are computed using the original, unbiased logits:
Because is excluded from the softmax denominator, the scale of reflects pure model-driven confidence rather than artificial routing corrections.
- Runtime Bias Update Rule: At the end of each training step , the bias term for each expert is adjusted directly based on its observed token load relative to the target average load :
where is a fixed step size (typically ), and denotes the signum function:
- If expert is underloaded (), increases, raising the probability that expert enters the top- in subsequent steps.
- If expert is overloaded (), decreases, throttling incoming token traffic.
Why Loss-Free Balancing Outperforms Auxiliary Losses
- Zero Gradient Interference: Because is updated as an external control variable rather than via autograd backpropagation, . The routing projection updates solely to minimize the primary language modeling cross-entropy .
- Deterministic Load Convergence: Unlike stochastic gradient descent which oscillates around auxiliary loss landscapes, the signum control loop acts as an integral feedback controller, maintaining expert utilization within of perfect parity throughout training.
- Elimination of Dropped Tokens: Perfect load parity allows distributed clusters to operate at with zero token drops, reducing memory allocation overhead while preserving maximum downstream reasoning performance.
9. Architectural Comparison of MoE Routing Paradigms
| Routing Mechanism | Gating Formulation | Load Balancing Method | Dropped Tokens | Gradient Interference | Primary Production Implementations | | :--- | :--- | :--- | :--- | :--- | :--- | | Noisy Top-k Gating | | Stochastic Gaussian noise exploration | Yes (High under load imbalance) | Minimal | Shazeer et al. (2017) | | Switch Routing | | Auxiliary Loss () | Yes (If ) | Moderate | Switch Transformer | | ST-MoE Gating | | Auxiliary Loss + Router Z-Loss | Yes (Low with capacity slack) | Moderate | ST-MoE, PaLM-2 MoE | | Expert-Choice | | Column-wise token selection | No (0% by construction) | None | Zhou et al. (2022) | | Fine-Grained + Shared | | Auxiliary Loss + Shared Invariant Routing | Yes (Bounded) | Low-to-Moderate | DeepSeekMoE, Qwen-MoE | | Auxiliary-Loss-Free | with unperturbed | Dynamic signum bias control loop () | No (Zero drop at ) | None (Zero) | DeepSeek-V3, DeepSeek-R1 |
10. Summary and Implementation Takeaways
The evolution of Mixture-of-Experts routing reflects a continuous push toward eliminating artificial constraints on the language modeling objective:
- Routing Collapse is an Inherent Dynamic of Softmax Selection: Naive gradient descent on sparse gating layers inevitably produces winner-take-all routing.
- Auxiliary Losses are a Compromise: While inner-product auxiliary penalties () force equal distribution, their gradient backpropagation actively conflicts with task performance, forcing a compromise between hardware efficiency and token perplexity.
- Z-Loss Resolves Precision Instability: Adding a squared log-partition regularizer prevents logit magnitude explosion, protecting FP16/BF16 training stability across deep layers.
- Shared Experts Prevent Redundant Parameter Allocation: Dedicated invariant experts absorb general language representations, allowing routed experts to focus entirely on orthogonal specialized domains.
- Dynamic Biases Decouple Load Control from Representation Learning: Shifting load balancing from loss-function gradient backpropagation to runtime feedback control (as in DeepSeek-V3) achieves hardware load balancing without degrading language modeling quality.
Sources
- Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. arXiv:1701.06538.
- Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Yadav, Y., ... & Chen, Z. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ICLR 2021.
- Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Journal of Machine Learning Research (JMLR).
- Zoph, B., Bello, I., Kumar, S., Du, N., Huang, D., Dean, J., ... & Fedus, W. (2022). ST-MoE: Designing Stable and Transferable Sparse Expert Models. arXiv:2202.08906.
- Zhou, Y., Lei, T., Liu, H., Du, N., Huang, D., Zhao, V. Y., ... & Le, Q. V. (2022). Mixture-of-Experts with Expert Choice Routing. NeurIPS 2022.
- Dai, D., Deng, C., Zhao, C., Xu, R. X., Gao, H., Chen, D. L., ... & Liang, W. (2024). DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. arXiv:2401.06066.
- Wang, P., et al. (2024). Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts. arXiv:2408.15664.



