Load Balancing › Load Metrics

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:

$$f_i = \frac{1}{N} \sum_{t=1}^{N} \mathbb{1}[\text{argmax}_j\, g_j^{(t)} = i]$$

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:

$$\text{LIF} = E \cdot \max_i f_i$$

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:

$$\text{CV} = \frac{\sigma(f)}{\mu(f)} = \frac{\sigma(f)}{1/E}$$

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:

$$\bar{g}_i = \frac{1}{N} \sum_{t=1}^{N} g_i^{(t)}$$

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:

$$H(\bar{g}) = -\sum_{i=1}^{E} \bar{g}_i \log \bar{g}_i$$

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.