MoE Load Balancing: Token Distribution & Expert Collapse

In the previous chapters, we explored how Mixture of Experts architectures use gating networks to route tokens to specialized expert networks. Top-k routing selects the most relevant experts for each token, allowing sparse computation while maintaining model capacity. By activating only a small fraction of the model's parameters for any given input, MoE architectures can scale model capacity beyond what would be computationally feasible with dense networks. A model with 64 experts and top-2 routing activates only about 3% of its expert parameters per token, yet retains access to all expert knowledge through the training process. This routing mechanism also introduces a challenge that can undermine the architecture's benefits.

The gating network can develop strong preferences for certain experts while ignoring others. This creates an imbalance that wastes model capacity, causes computational inefficiencies during distributed training, and in extreme cases leads to complete model collapse. This failure mode has derailed training runs, caused wasted compute, and motivated research on keeping experts busy. To train MoE models reliably, you need to understand why they break.

Load balancing addresses these problems by so all experts receive a roughly equal share of training tokens. Without explicit balancing mechanisms, MoE models consistently converge to pathological states where only a handful of experts process the large majority of tokens. The remaining experts sit idle, consuming memory and adding model parameters that contribute nothing to the output. This is one of the central paradoxes of MoE training: the architecture is designed to provide more capacity, yet without careful attention to routing dynamics, you often end up with less effective capacity than a simpler dense model.

Understanding why this happens and how to measure it is needed before exploring solutions like auxiliary losses in subsequent chapters. This chapter builds the diagnostic foundation: we examine the mechanics of imbalance, characterize failure modes, and develop quantitative tools to detect problems before they become catastrophic. Think of this chapter as learning to read an MoE model's diagnostic signals. Before you can prescribe treatment, you need to know what healthy behavior looks like and how to recognize failures.

In the next chapters, we'll cover auxiliary balancing losses and router z-loss, which are the primary tools practitioners use to actively correct these imbalances during training. For now, the goal is deep understanding of the problem itself.

When a gating network routes tokens to experts, nothing inherently prevents it from developing a preference for certain experts over others. In fact, the opposite is true: the training dynamics of MoE systems naturally push toward imbalanced routing. This tendency emerges not from any flaw in the architecture design but from the basic way that learning systems respond to feedback signals.

Think of the gating network's job as similar to a hiring manager who must allocate work across a team of consultants. At the start of a project, the manager has no basis for preferring one consultant over another, so work is distributed somewhat evenly. But as the project progresses, certain consultants demonstrate strong performance on early tasks. Rationally, the manager routes more important work to these consultants. The more they work, the more they develop their skills and demonstrate capability, reinforcing the manager's preference. Meanwhile, the consultants who receive less work have fewer opportunities to demonstrate their abilities, and over time their skills may even atrophy. By the end of the project, a few consultants are overworked while others sit idle, even though the original team was assembled to handle a diverse workload. The MoE gating network is this manager, and the experts are the consultants.

To understand why imbalance occurs at a mechanistic level, consider what happens during the early stages of training. If expert $E_3$ happens to perform slightly better than its peers on a few early batches, perhaps due to nothing more than favorable random initialization, the gating network learns to send more tokens to $E_3$. This makes sense from an optimization perspective: the gating network is simply learning to route tokens to whichever expert currently produces the best outputs. However, this reasonable local decision creates a problematic global pattern.

With more tokens flowing to $E_3$, this expert receives more gradient updates and accumulates more training signal. These additional updates allow $E_3$ to improve further, which in turn reinforces the gating network's preference for routing to this expert. Meanwhile, experts that receive fewer tokens get less training signal and fall further behind. They have fewer opportunities to learn from data, fewer gradient updates to refine their weights, and consequently less capability to offer when tokens are occasionally routed their way. This disparity creates a positive feedback loop that amplifies initial imbalances over time.

The key insight is that the MoE system contains a hidden instability. The gating network's training objective, to route tokens well, creates incentives that undermine the system's own capacity utilization. Every local improvement in routing quality slightly worsens the global distribution of learning. This is why load balancing cannot be an afterthought: it must be built into the training procedure from the start.

The problem manifests at two distinct but interconnected levels. At the token level, some experts process disproportionately many tokens within each batch, meaning that the computational load is not evenly distributed across the available expert networks. At the training level, some experts receive far more gradient updates over the course of training, meaning that certain experts learn extensively while others remain undertrained. Both types of imbalance waste model capacity, and together they create a situation where the theoretical advantages of MoE architectures cannot be realized in practice. The model pays the memory and complexity costs of having many experts while only benefiting from a few.

An 8-expert MoE layer with severe imbalance might route 60% of tokens to two experts while the remaining six experts share only 40% of tokens. In the extreme, several experts might receive essentially no tokens at all, sitting idle while a small subset of their peers handle all the processing workload. From the outside, the model looks like an 8-expert system. From the inside, it is a 2-expert system wearing an expensive costume.

This imbalanced state creates two distinct but equally serious problems that undermine the value of the MoE architecture:

• Wasted parameters: Underutilized experts contribute little to model predictions despite consuming memory. A model with 8 experts where only 3 are active effectively has the capacity of a 3-expert model with extra overhead. The memory budget allocated to store the weights of unused experts could have been spent on additional layers, larger hidden dimensions, or other architectural improvements that would benefit the model's performance.

• Training bottlenecks: In distributed settings where each expert resides on a different accelerator, load imbalance creates stragglers. The accelerators hosting popular experts finish processing their tokens while others sit idle, then must wait for synchronization. This waiting time represents pure waste: expensive hardware consuming power without creating useful computation.

The computational inefficiency is particularly acute in expert parallelism, where each expert is placed on a separate device. Consider a scenario where expert $E_1$ receives 10 times more tokens than expert $E_7$. In this case, the device hosting $E_1$ performs 10 times more computation while the device hosting $E_7$ mostly waits. Since training requires synchronization across all devices at regular intervals, the entire system is bottlenecked by the slowest device. Paradoxically, the slowest device in this context is not the one with the least capable hardware but rather the one hosting the most overloaded expert.

There is also an insidious interaction between load imbalance and batch size. As imbalance grows, the effective batch size for each expert diverges. The overloaded experts process more examples and receive stronger, lower-variance gradient estimates. The underloaded experts process fewer examples and receive higher-variance, noisier gradients. The resulting training dynamics are uneven: some parts of the model converge quickly and reliably while others wander due to gradient noise. This uneven convergence makes the model harder to tune and the final result harder to predict.

The most severe manifestation of load imbalance is expert collapse, a failure mode where the model converges to using only one or two experts for nearly all tokens. It breaks the sparse-computation promise of MoE. When collapse occurs, the model has effectively reverted to a dense architecture, paying the memory costs of multiple experts while receiving the computational benefits of only one or two. The ability to scale model capacity without proportionally scaling computation is lost.

Expert collapse is one of those failure modes that seems surprising at first but makes perfect sense once you understand the underlying dynamics. The architecture is designed to be flexible, and this flexibility, without constraints, allows it to find a degenerate solution where the gating network simply routes everything to the best-performing expert and ignores the rest. From the optimizer's perspective, this is a perfectly valid strategy: if one expert can handle all the tokens adequately, why maintain eight? The problem is that this strategy violates the whole point of having eight experts in the first place.

Expert collapse occurs when the rich-get-richer dynamics run to their logical extreme. The process typically unfolds through a predictable sequence of stages, each building on the previous one:

1. Early training: The gating network's random initialization creates slight expert preferences. These initial biases are typically small and might seem insignificant, but they provide the seed from which imbalance can grow.
2. Divergence: Preferred experts receive more tokens and more gradients, improving faster. The gap between experts begins to widen, though it may still be subtle enough to escape notice in standard training metrics.
3. Reinforcement: The gating network learns to rely more heavily on better experts. At this stage, the feedback loop is firmly established and accelerating. The router's preferences become increasingly pronounced.
4. Saturation: A small subset of experts handles nearly all tokens. The model has effectively reorganized itself around a handful of dominant experts, with others receiving only occasional, often random, token assignments.
5. Collapse: Non-preferred experts stop receiving training signal entirely. The process reaches its endpoint, where certain experts are completely excluded from the model's functioning.

Once collapsed, recovery is nearly impossible through normal training. The ignored experts become stale: their weights remain at initialization values or drift randomly without meaningful gradients. The gating network has learned to avoid them, so they never get the chance to improve. This creates a stable but pathological equilibrium where the model is stuck with reduced capacity and no path back to balanced utilization.

Think of expert collapse as analogous to a market monopoly. Once a dominant player captures enough market share, the remaining competitors cannot attract enough business to invest in improvements, which causes them to fall further behind, which causes them to lose even more business. The end state, where one player controls everything, is self-sustaining even if it is inefficient for the overall system. Breaking out of this equilibrium requires external intervention, just as MoE training requires explicit balancing mechanisms.

Collapse can be detected by monitoring expert utilization during training. A healthy MoE layer shows all experts receiving tokens in roughly equal proportions, with natural fluctuations based on the content of each batch. A collapsing layer shows increasing concentration in a small number of experts, with the trend persisting across batches rather than averaging out over time.

The warning signs include:

• Monotonically increasing utilization for a subset of experts across consecutive batches
• Near-zero routing probabilities for some experts that persist and worsen over time
• Decreasing entropy in the router's softmax outputs, measured as a trend rather than momentary fluctuation
• Sharp drops in validation performance as the model loses access to specialized expert knowledge
• Increasing gradient norms for popular experts and decreasing norms for unpopular ones

The important distinction is between natural variation and systematic trend. In any given batch, some experts will naturally receive more tokens than others simply because the batch content favors certain topics. This per-batch variation is healthy and expected. Collapse is characterized by trends that persist and amplify across many batches. Monitoring tools should compute rolling averages over many steps to distinguish signal from noise.

Catching collapse early is important because intervention becomes harder as the process advances. In the divergence stage, simple adjustments to the auxiliary loss coefficient or learning rate may be sufficient. By the saturation stage, the model may require restarting from an earlier checkpoint with modified hyperparameters. At full collapse, recovery through continued training is practically impossible.

The softmax operation in the gating network contributes to collapse in a basic way that makes this failure mode particularly difficult to avoid. Recall from the gating networks chapter that router scores pass through a softmax to produce routing probabilities.

To understand this quantitatively, consider how softmax transforms a set of scores. The softmax function takes a vector of real-valued scores and converts them into a probability distribution, where all values are positive and sum to one:

where:

• $g_i$: the routing probability assigned to expert $i$
• $s_i$: the raw score (logit) for expert $i$
• $E$: the total number of experts
• $e^{s_i}$: the exponential function applied to the score, so a positive value
• $\sum_{j=1}^{E} e^{s_j}$: the normalizing sum over all $E$ experts

Why does this formula make sense as a routing mechanism? Notice that the exponential function is always positive (no matter how negative the score), ensures that all outputs are non-negative, and that dividing by the sum guarantees outputs sum to one. It is a clean way to turn arbitrary real-valued scores into valid probabilities.

This formula reveals why collapse is self-reinforcing. The exponential function at the heart of softmax has a powerful sharpening effect on probability distributions. Small differences in raw scores $s_i$ get amplified by the exponential, and larger differences get amplified even more. To understand this intuitively, consider that if one score increases by just 1 unit, its contribution to the numerator increases by a factor of approximately 2.7 (since $e^1 \approx 2.718$). This multiplicative effect means that even modest advantages in raw scores translate into substantial advantages in routing probability.

If expert $E_1$ consistently has slightly higher scores than others, the softmax concentrates probability mass on $E_1$, reducing training signal to other experts. As the gap grows over successive training steps, the softmax's sharpening effect accelerates collapse. What starts as a small preference can rapidly snowball into complete dominance. This is why monitoring metrics like router entropy, which we discuss shortly, provides valuable early warning signals before collapse becomes irreversible.

To address load imbalance, we first need to quantify it. Measuring imbalance precisely allows us to detect problems early, compare different balancing strategies, and tune hyperparameters effectively. Several metrics capture different aspects of expert utilization, each giving a distinct perspective on the health of the routing system.

Think of these metrics as instruments in a diagnostic toolkit. Just as a doctor might monitor heart rate, blood pressure, and oxygen saturation simultaneously because each captures something different about patient health, we monitor multiple load metrics simultaneously because each reveals a different facet of routing behavior. No single metric tells the complete story, but together they provide a complete picture.

The most straightforward diagnostic question is: are all experts receiving tokens? But this alone is insufficient. An expert receiving 2% of tokens instead of 12.5% (its fair share in an 8-expert system) might not appear completely idle, yet the 6x shortfall still causes significant capability degradation. We need metrics that make these disparities numerically explicit and comparable across different model configurations.

The most direct and intuitive measure is the fraction of tokens routed to each expert within a batch. This metric simply counts how many tokens each expert receives and divides by the total, giving us a clear picture of the current workload distribution.

Before writing the formula, consider what we want to compute: for each expert, we want to count how many tokens chose that expert as their top choice, then normalize by the total number of tokens. This gives a number between 0 and 1 that represents the expert's share of the workload. For a batch of $N$ tokens with top-1 routing, the token fraction for expert $i$ is:

where:

• $f_i$: the fraction of the batch assigned to expert $i$
• $N$: the total number of tokens in the batch
• $t$: the index variable iterating over all tokens
• $\mathbb{1}[\cdot]$: the indicator function, which is 1 if the condition inside is true and 0 otherwise
• $g_j^{(t)}$: the routing probability assigned to expert $j$ for token $t$
• $\text{argmax}_j\, g_j^{(t)}$: the index of the expert with the highest probability for token $t$
• $i$: the specific expert we are measuring

The formula works by iterating through every token in the batch and checking whether that token was routed to expert $i$. The indicator function returns 1 when token $t$'s highest-probability expert matches $i$, and 0 otherwise. Summing these indicators and dividing by $N$ gives us the fraction of all tokens that expert $i$ processed.

Why does this formula make sense? The indicator function enforces the hard routing decision: a token either went to expert $i$ (contributes 1) or it did not (contributes 0). This discrete counting reflects how expert capacity is consumed during the forward pass. You cannot half-route a token; it goes to exactly one expert (for top-1 routing). The sum and division are standard fraction computation.

Perfect balance means $f_i = \frac{1}{E}$ for all $E$ experts. In practice, some deviation is expected and even desirable, since not all tokens should require the same expert. Natural language contains diverse phenomena, and different token types may need different processing. However, severe skew indicates problems that will compound over training.

While token fractions tell us how many tokens each expert received, interpreting these raw numbers requires context. Is a token fraction of 0.15 for one expert concerning in an 8-expert system? What about in a 16-expert system? The load imbalance factor provides a normalized measure that answers these questions by measuring how far from uniform the distribution is.

The key idea is to compare the actual maximum load to the theoretical ideal load. In a perfectly balanced system, the most-loaded expert would receive exactly $1/E$ of all tokens. The load imbalance factor tells us how many times larger the actual maximum is compared to this ideal:

where:

• $\text{LIF}$: the Load Imbalance Factor
• $E$: the total number of experts
• $\max_i f_i$: the highest token fraction observed among all experts
• $f_i$: the fraction of tokens assigned to expert $i$

This metric compares the maximum utilization to the ideal uniform utilization ($1/E$). The intuition behind this formula is straightforward: if all experts received exactly their fair share, the maximum fraction would be $1/E$, and multiplying by $E$ would give exactly 1. Any deviation from uniformity pushes this value higher.

Why does this formula make sense? Notice that $E \cdot (1/E) = 1$, so perfect balance yields $\text{LIF} = 1$. If one expert receives all tokens ($\max_i f_i = 1$), then $\text{LIF} = E$. This gives us a natural scale from 1 (perfect) to $E$ (complete collapse). The key insight is that the formula normalizes by the number of experts, so a LIF of 2.0 always means the same thing regardless of how many experts the model has: the most popular expert is doing twice its fair share.

With perfect balance ($f_i = \frac{1}{E}$ for all $i$), the load imbalance factor equals 1. If all tokens go to a single expert ($f_i = 1$ for one expert), the factor equals $E$. Values above 1 indicate imbalance, with higher values being worse. This scaling makes the metric interpretable across different numbers of experts: a load imbalance factor of 2 always means the most-used expert receives twice its fair share, regardless of whether the system has 4 experts or 64.

Another useful metric is the coefficient of variation of expert loads, which captures how spread out the utilization values are relative to their expected value. While the load imbalance factor focuses on the single worst-case expert, the coefficient of variation considers the entire distribution.

We first compute the standard deviation of all expert token fractions and then normalize by the mean fraction:

where:

• $\text{CV}$: the Coefficient of Variation
• $\sigma(f)$: the standard deviation of the token fractions across all experts
• $\mu(f)$: the mean token fraction (equal to $1/E$ since fractions sum to 1)
• $E$: the number of experts

The coefficient of variation differs from the load imbalance factor in what it measures. While the load imbalance factor focuses on the single most overloaded expert, the coefficient of variation considers the entire distribution. A coefficient of variation of 0 indicates perfect balance, where all experts receive identical token counts. Higher values indicate greater dispersion in utilization.

Why does this formula make sense? Notice that dividing the standard deviation by the mean (rather than reporting the raw standard deviation) makes the metric scale-invariant. A system with 8 experts has a mean fraction of 0.125, while a system with 64 experts has a mean fraction of 0.0156. The coefficient of variation correctly normalizes for this difference, so comparisons across different model configurations are meaningful.

This metric is particularly useful for detecting cases where multiple experts are both over- and under-utilized, even if no single expert is dramatically overloaded. Consider two scenarios: in the first, one expert has 40% of tokens and all others share the remainder equally; in the second, half the experts have 20% each and the other half have 5% each. The load imbalance factor treats these differently (2x vs 1.6x), but the coefficient of variation captures the more diffuse imbalance in the second scenario that the max-based metric might understate.

Rather than measuring hard assignments, which can mask the router's underlying preferences, we can also examine the soft probabilities from the router. Hard assignments reflect what happened during the forward pass, but the underlying probabilities indicate what the router prefers. These are related but distinct: a router might assign equal numbers of tokens to two experts (equal hard assignments) while consistently preferring one over the other in the soft probability values.

The mean routing probability for expert $i$ across a batch is computed by averaging the softmax output for that expert position across all tokens:

where:

• $\bar{g}_i$: the average probability assigned to expert $i$ across the batch
• $N$: the number of tokens in the batch
• $g_i^{(t)}$: the probability assigned to expert $i$ for token $t$

This metric reveals information that token fractions alone cannot. Even if tokens are routed via top-k selection (which makes discrete choices), the underlying probabilities reveal the router's preferences before discretization. An expert might receive the same number of tokens as its peers in terms of final assignments but might consistently be the second choice rather than the first, showing that it is at risk of falling behind.

We can quantify the dispersion of these soft probabilities using the entropy of the mean routing probabilities. Entropy measures how concentrated or spread out a probability distribution is. A uniform distribution has maximum entropy, while a distribution that puts all mass on one outcome has zero entropy. We compute it as:

where:

• $H(\bar{g})$: the entropy of the average routing distribution
• $E$: the total number of experts
• $\bar{g}_i$: the average probability assigned to expert $i$
• $\log$: the natural logarithm

Why does this formula make sense? Notice the negative sign: the formula computes $-\sum p \log p$, the standard Shannon entropy. When probabilities are uniform ($\bar{g}_i = 1/E$ for all $i$), this achieves its maximum value of $\log E$. When all probability mass concentrates on a single expert ($\bar{g}_1 = 1$, all others 0), the entropy equals 0 (since $1 \cdot \log 1 = 0$ and the other terms vanish). Higher entropy means more balanced routing; lower entropy signals concentration and incipient collapse.

Maximum entropy occurs when all experts have equal average probability ($\bar{g}_i = 1/E$ for all $i$), corresponding to a uniform distribution. As the router's preferences sharpen toward collapse, entropy decreases toward zero. Entropy is a sensitive early warning metric because it captures the complete distribution rather than only the maximum and begins to decrease as the distribution departs from uniformity, often before token-fraction metrics show alarming values.

Beyond computational efficiency, balanced routing directly affects model quality. The relationship between load balance and model performance involves more than keeping hardware busy. Balance determines whether the MoE architecture can use its distributed capacity as intended.

The core promise of MoE is that different experts can specialize in different aspects of language. One expert might become skilled at handling mathematical notation, another at processing idiomatic expressions, another at managing long-range syntactic dependencies. This specialization emerges during training as each expert is repeatedly exposed to the token types where it performs best and receives gradients that sharpen those capabilities. Load imbalance disrupts this specialization process at every stage, from the initial emergence of expert differences to the maintenance of those differences at scale.

Each expert in an MoE layer contains independent parameters that can specialize in different aspects of the input. MoE gains capacity through distributed specialization: rather than a single network trying to handle all patterns, multiple networks can each become expert in their own domain. An 8-expert layer with 64M parameters per expert has 512M total parameters. If only 3 experts are used, the effective capacity drops to 192M parameters while memory consumption remains at 512M. Balanced routing uses the full parameter budget and achieves the capacity efficiency that motivates MoE designs.

This capacity efficiency argument has concrete implications for model scaling. One of the main reasons practitioners adopt MoE architectures is the favorable scaling law: as you add more experts, model capacity grows while inference computation per token remains roughly constant (for fixed top-k). But this scaling argument assumes all experts are being used. If imbalance causes half the experts to be idle, you are paying the full capacity cost while getting only half the capacity benefit. The scaling law breaks down precisely when you need it most.

Experts that receive few tokens get noisy gradient estimates. The mathematics of stochastic gradient descent relies on averaging over many samples to get accurate estimates of the true gradient direction. In the extreme, an expert that sees only 10 tokens per batch has high-variance gradients compared to one that sees 1,000 tokens. This variance makes optimization unstable for underutilized experts, preventing them from learning effectively even when they do receive tokens.

The resulting expert weights may oscillate rather than converge, never settling into useful representations. This oscillation happens because each small batch of tokens provides an inaccurate estimate of the true gradient direction. With enough samples, these errors average out; with too few samples, the expert's weights bounce between directions suggested by different small batches, making no net progress. The expert appears to be learning, because its weights are changing, but the changes are random rather than purposeful.

There is a compounding effect here: because noisy gradients produce poor expert quality, the router correctly learns to avoid poorly-performing experts, which reduces their token count further, which makes their gradients noisier, which worsens their performance. This is the rich-get-richer dynamic viewed from the gradient quality perspective.

With balanced routing, experts are forced to handle diverse token types. This encouragement of diversity leads to different experts specializing in different linguistic phenomena: one might become skilled at handling numerical expressions, another at processing named entities, and still another at understanding syntactic structures. If routing is imbalanced, overutilized experts become generalists (handling everything), while underutilized experts never develop meaningful specializations.

The model loses the specialization benefits that make MoE architectures valuable. At the extreme, a model with collapsed routing might perform worse than a comparable dense model: it has the same number of active parameters per token but has wasted its capacity on redundant generalist experts rather than developing the diverse specializations that MoE enables. The model ends up paying the memory cost of many experts while getting the diversity benefit of only one.

Specialization diversity also matters for generalization. A model where experts have different capabilities is more reliable under distribution shift: different test domains activate different experts that use their specialized knowledge. A model with collapsed routing has all its eggs in one basket. When that single dominant expert encounters unfamiliar input patterns, there are no alternative experts with complementary knowledge to fall back on.

Severe imbalance creates training instabilities that manifest in multiple ways. As expert utilization shifts, the effective model capacity changes mid-training. The model might behave differently from one epoch to the next not because it is learning but because different experts are being used. Sharp routing transitions, where the gating network suddenly starts preferring a different expert, cause loss spikes and can disrupt convergence. These instabilities make training progress unpredictable and can require extensive hyperparameter tuning to mitigate.

The instabilities also interact badly with learning rate schedules. Standard schedules like cosine annealing or linear decay are designed assuming a stable training trajectory. When expert utilization shifts cause loss spikes, these schedules may be in the wrong phase: the learning rate might be too high during a spike (amplifying instability) or too low during a recovery (slowing progress). Practitioners training MoE models often report needing to use more conservative learning rate schedules than they would for equivalent dense models, which means slower overall training.

Let's implement functions to compute load balance metrics and visualize expert utilization across training. The goal is to build the diagnostic toolkit you'll use when training and debugging your own MoE models. We'll start with the foundational measurements and then build up to simulations that show how imbalance evolves during training.

The output confirms we have 2,048 tokens in total (16 sequences $\times$ 128 length). The routing probability tensor shape matches our expectation, containing a probability distribution over the 8 experts for every token in the batch.

Let's compute the fraction of tokens routed to each expert. The function below implements the token fraction formula directly: for each expert, count the tokens assigned to it and divide by the total.

The load imbalance factor above 1.0 indicates that the most popular expert receives more than its fair share of tokens. The normalized entropy below 1.0 shows that the distribution is not uniform. These numbers represent a model that has not even started training, and the biases we introduced are small compared to what develops during actual training runs. This gives you a sense of how aggressive the rich-get-richer dynamics can be.

To see how imbalance develops during training, let's simulate the rich-get-richer dynamics by iteratively biasing the gating network toward popular experts. This simulation compresses the feedback loop that plays out over real training runs into a controlled demonstration.

The simulation demonstrates how quickly expert utilization diverges from uniform. After only 100 steps, the load imbalance factor has risen substantially, with some experts receiving several times their fair share of tokens while others receive almost none. In a real training run, this feedback loop operates over hundreds of thousands of gradient steps, giving it ample time to reach pathological extremes.

Let's run a longer simulation with stronger reinforcement to observe expert collapse. We increase both the number of steps and the reinforcement strength to show what happens when no balancing mechanism intervenes.

The stacked area chart shows expert collapse in action. By the end of training, nearly all tokens are routed to just one or two experts. The other experts, despite consuming memory and having trainable parameters, contribute nothing to model outputs. Notice how the collapse happens in stages: initially, some experts grow while others shrink, then the dominant experts gradually squeeze out all competitors until the final collapsed state is reached.

Router entropy provides another view into load balance. Let's examine how the softmax output distribution changes as collapse progresses. Entropy is particularly valuable because it begins declining earlier than other metrics. This provides a leading indicator of collapse before it becomes obvious in the token fraction plots.

Looking at all three metrics together, we can see how they provide complementary information. The Load Imbalance Factor rises sharply once the dominant expert clearly separates from the pack, which makes it good for detecting moderate to severe imbalance. The Coefficient of Variation is sensitive to diffuse imbalance where many experts deviate from their ideal load, even if no single expert is dramatically dominant. Normalized Entropy provides the earliest warning signal, declining smoothly from the very first steps as the distribution begins to concentrate.

Load imbalance has particularly severe consequences when experts are distributed across devices. In expert parallelism, each expert resides on a separate accelerator, and the system's throughput is constrained by the slowest device. This is a basic result from parallel computing theory known as Amdahl's Law: the speedup from parallelization is limited by the fraction of work that cannot be parallelized, or in this case, by the time required for the most overloaded device. Let's quantify the efficiency loss that results from uneven workload distribution.

Think of the distributed training scenario as a team of runners in a relay race where all runners must cross the finish line before the next leg begins. If some runners are much faster than others, the fast runners finish early and wait while the slow runners complete their legs. The total time is determined by the slowest runner, not the average. Load imbalance in expert parallelism creates exactly this situation: the device hosting the overloaded expert is the slow runner, and all other devices wait for it to finish before training can proceed.

The relationship is stark: efficiency is inversely proportional to the load imbalance factor. When the most-loaded expert receives twice its fair share of tokens (LIF = 2), parallel efficiency drops to 50%. When one expert receives all tokens (LIF = $E$), efficiency approaches $\frac{1}{E}$ since only one device does useful work while the others wait.

Consider training a large MoE model with 64 experts distributed across 64 GPUs. At perfect balance, all 64 GPUs compute in parallel, giving approximately 64x speedup over a single GPU. If load collapses such that one expert handles 50% of tokens (LIF = 32), efficiency drops to about 3%, meaning that 62 out of 64 GPUs spend most of their time waiting. The effective throughput is barely better than what a single GPU could achieve. Load balancing is therefore a practical necessity for serious MoE training at scale.

Let's walk through a concrete numerical example to solidify your understanding of how these metrics evolve together. We'll trace the values at three snapshots through a collapse simulation: the beginning, the middle, and the end.

Suppose we start with an 8-expert system at step 0, where all experts receive exactly uniform token fractions of 0.125. At this point, the metrics are:

• Token fractions: each $f_i = 0.125$
• LIF: $8 \times 0.125 = 1.0$ (perfect balance)
• CV: $\sigma(f) = 0$, so $\text{CV} = 0$ (no variation)
• Normalized entropy: $-8 \times (0.125 \times \log 0.125) / \log 8 = 1.0$ (maximum)

Now suppose we reach step 50, where the feedback loop has been running with moderate strength. Expert 0 now receives 0.25 of tokens, experts 1 and 2 receive 0.15 each, experts 3-6 receive 0.10 each, and expert 7 receives 0.05. Let's compute:

• LIF: $8 \times 0.25 = 2.0$ (most loaded expert has twice its fair share)
• Mean: $1/8 = 0.125$, Standard deviation $\approx 0.063$, so $\text{CV} \approx 0.063 / 0.125 = 0.504$
• Entropy requires summing $-f_i \log f_i$ for each expert; this gives a normalized entropy around 0.91, showing the distribution is noticeably but not critically skewed

Finally at step 200, after strong reinforcement, expert 0 receives 0.70 of all tokens, expert 1 receives 0.25, and the remaining 6 experts share 0.05 total (approximately 0.008 each). The metrics now read:

• LIF: $8 \times 0.70 = 5.6$ (expert 0 receives 5.6 times its fair share)
• CV: very high, around 2.0 or above, which reflects massive spread
• Normalized entropy: close to 0.3, far below the maximum, showing highly concentrated routing

This progression illustrates the value of monitoring all three metrics simultaneously. The LIF gives you the severity of the worst-case situation. The CV reveals whether the problem is concentrated (one dominant expert) or diffuse (many experts off their ideal load). The entropy provides the earliest warning, declining first and most sensitively. In practice, you would set alerting thresholds: for example, flag a potential problem when LIF exceeds 1.5 or normalized entropy falls below 0.85, and intervene when LIF exceeds 2.0 or entropy falls below 0.7.

One approach to limiting imbalance is to impose a hard capacity constraint on each expert. Rather than letting the router to send arbitrarily many tokens to a popular expert, we define a maximum limit that caps how many tokens any single expert can process. This constraint prevents the most extreme forms of imbalance by rejecting tokens that would overflow an expert's capacity.

Think of capacity constraints as a fire code for experts. Just as a concert venue has a maximum occupancy limit that prevents dangerous overcrowding (even if more people want to attend), each expert has a maximum token capacity that prevents it from being overwhelmed. Tokens that arrive at a full expert are turned away, just as latecomers to a sold-out venue are turned away at the door.

We define a maximum capacity limit for each expert using the following formula. Given the total number of tokens $N$ and the number of experts $E$, the ideal load per expert is $N/E$. The capacity factor $C$ controls how much we allow individual experts to exceed this ideal:

where:

• $K_{\text{max}}$: the maximum number of tokens an expert is allowed to process
• $C$: the capacity factor (typically $>1.0$), which determines the buffer size
• $N$: the total number of tokens in the batch
• $E$: the number of experts
• $\lfloor \cdot \rfloor$: the floor function, so an integer result

This formula sets the limit to a multiple $C$ of the ideal uniform load ($N/E$), letting for natural variance while capping extreme imbalance. The capacity factor $C$ provides flexibility: a value of 1.0 would enforce strict uniformity, while higher values like 1.25 or 1.5 allow experts to handle somewhat more than their fair share before hitting the limit. The floor function ensures we get a whole number of tokens, since partial token assignments are not meaningful.

Why does this formula make sense? Notice that when $C = 1$, $K_{\text{max}} = N/E$, exactly the uniform allocation. When $C = 2$, experts can handle double their fair share before tokens start being dropped. The parameter $C$ lets you tune the tradeoff between strict enforcement (low $C$, more dropped tokens, better balance) and lenient enforcement (high $C$, fewer dropped tokens, worse balance). Most implementations use $C$ between 1.1 and 1.5 to allow modest overloading without creating severe imbalance.

Capacity constraints prevent catastrophic imbalance by rejecting tokens that would overflow an expert's capacity. The mechanism is effective at bounding the maximum load any single expert can experience, which directly bounds the Load Imbalance Factor: with a capacity factor of 1.25, the LIF can never exceed 1.25. However, this approach has a significant downside: dropped tokens receive no expert processing and must either be handled by a shared residual pathway or simply contribute less to the model's output.

The dropped-token problem is significant. Each dropped token represents information that the model fails to process correctly. In a language-modeling context, this might mean the model makes worse predictions for the next token because some context tokens were not processed by their appropriate expert. At high drop rates, this can degrade model quality.

The next chapter explores auxiliary losses that encourage balance during training rather than enforcing it through hard constraints. These losses add a penalty term to the training objective that discourages the router from developing strong preferences for specific experts. The key advantage over capacity constraints is that auxiliary losses are proactive rather than reactive: they prevent imbalance from developing in the first place, so tokens are never dropped.

Understanding the key parameters that govern load balance helps you configure MoE systems effectively and diagnose problems when they arise. These parameters interact in ways that require careful consideration:

• num_experts: The total number of experts in the MoE layer. More experts increase model capacity but make balancing harder because the probability mass must be spread across more options. With 8 experts, each expert's fair share is 12.5%; with 64 experts, it is only 1.56%. Small absolute deviations become more significant in proportional terms as the number of experts grows.

• capacity_factor: A scalar multiplier (typically 1.0 to 1.25) determining the maximum tokens an expert can process relative to uniform load. Lower values enforce stricter balance at the cost of more dropped tokens. Higher values accept more imbalance to reduce token dropping. The right value depends on the severity of your load imbalance and your tolerance for dropped tokens.

• reinforcement_strength: In our simulation, this parameter models the magnitude of the feedback loop. In real training, this corresponds to how strongly the gating network's gradients reinforce existing preferences, which depends on learning rate, batch size, and the auxiliary loss coefficient. Higher reinforcement leads to faster collapse; higher auxiliary loss coefficients counteract this by penalizing concentration.

Choosing these parameters well requires monitoring the load balance metrics we've discussed. Start by establishing baselines during initial training runs, set alerting thresholds based on those baselines, and tune the capacity factor and auxiliary loss coefficient until the metrics stay within acceptable ranges throughout training.

Load balancing in MoE models involves basic tensions that do not have clean resolutions. Understanding them supports principled tradeoffs rather than blindly optimizing metrics.

The first tension is between enforced balance and natural specialization. Perfect balance, where every expert receives exactly the same number of tokens, may not be optimal for model quality. Different token types require different processing, and some semantic categories may be more prevalent in the data than others. Forcing exact uniformity could degrade performance by preventing natural specialization patterns. Consider a model trained on a mix of English and code: code tokens might cluster toward certain experts that develop numerical or syntactic processing capabilities. Forcing those experts to process equal amounts of natural-language text and code could prevent them from fully specializing. The goal should be approximate balance that prevents pathological collapse rather than strict uniformity.

The second tension involves measurement stability. The metrics we've explored have limitations. Token fractions measured per batch can fluctuate substantially; an expert might appear underutilized in one batch but be needed for a specific topic that appears in the next. Longer-term averaging provides more stable estimates but responds slowly to shifts in expert utility. Entropy-based metrics capture the router's confidence distribution but do not distinguish between an expert appropriately confident on relevant tokens and one that is inappropriately dominant. No single metric tells the complete story, and over-optimizing any single metric can introduce distortions.

The third tension is between proactive and reactive balancing. Capacity constraints (reactive) prevent severe imbalance from causing computational problems but introduce token dropping, which degrades model quality. Auxiliary losses (proactive, covered in the next chapter) prevent imbalance from developing but add a training signal that may conflict with the primary task objective. Getting the balance right requires careful calibration: too little auxiliary loss and you still get collapse; too much and the auxiliary signal overwhelms the task signal, degrading the model's primary task performance.

Despite these limitations, load balancing is necessary for practical MoE deployment. Without it, the promised efficiency gains of sparse computation evaporate as models collapse to effectively dense computation through a small subset of experts. The training instabilities from severe imbalance, including gradient variance, capacity oscillation, and convergence failures, make unbalanced MoE models difficult to train at scale. Every major MoE system in production (Switch Transformer, GLaM, Mixtral, DeepSeek-MoE) uses some form of load balancing, which is the strongest possible endorsement for its necessity. These systems collectively represent trillions of tokens of training experience and billions of dollars of compute, and their builders independently concluded that balancing mechanisms are non-negotiable.

The failure to balance has real consequences beyond academic embarrassment. Training runs that collapse mid-way through are a significant waste of resources. For large models trained on expensive clusters, a single failed run due to expert collapse might cost hundreds of thousands of dollars and weeks of time. The investment in understanding and implementing proper load balancing is repaid many times over.

The next two chapters address load balancing through loss function design. Auxiliary balancing losses add a penalty term that encourages the router to spread tokens across experts. Router z-loss specifically targets the softmax's tendency to sharpen into collapse-prone distributions. Together with capacity constraints, these mechanisms enable stable training of MoE models with balanced expert utilization.

Load balancing ensures that all experts in a Mixture of Experts model receive meaningful training signal and contribute to model predictions. Without explicit balancing mechanisms, MoE training exhibits rich-get-richer dynamics where initially preferred experts attract more tokens, receive more gradients, and become even more preferred. This positive feedback loop leads to expert collapse, a state where only one or two experts handle nearly all tokens and the architecture's capacity advantage is entirely lost.

Key load balance metrics include the load imbalance factor (measuring deviation from uniform allocation as a multiple of the ideal), coefficient of variation (quantifying spread in token fractions across all experts), and router entropy (capturing the concentration of routing probability mass. This provides the earliest warning signal). These metrics are complementary: LIF catches severe single-expert dominance, CV detects diffuse multi-expert imbalance, and entropy monitors the router's internal preferences before they manifest as hard routing changes. Monitoring all three gives the most complete picture of routing health.

The consequences of imbalance extend beyond wasted model capacity. In distributed training with expert parallelism, load imbalance creates stragglers that destroy parallel efficiency. A load imbalance factor of 2 halves throughput, and collapsed models may run at only $\frac{1}{E}$ efficiency despite having $E$ experts. This makes load imbalance a model quality problem and a basic obstacle to economical training at scale.

Capacity constraints provide one mechanism for limiting imbalance by rejecting tokens that would overflow an expert's quota. The capacity limit $K_{\text{max}} = \lfloor C \cdot N/E \rfloor$ sets a soft ceiling that scales naturally with batch size and the number of experts. However, this approach drops tokens rather than preventing the router from developing imbalanced preferences in the first place. The auxiliary balancing loss and router z-loss covered in the next chapters address this by shaping the training objective itself, rewarding uniform routing and penalizing the sharp softmax outputs that accelerate collapse.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about load balancing in Mixture of Experts models.