Auxiliary Balancing Loss › Load Balancing Loss Formulation

The Switch Transformer paper (Fedus et al., 2021) introduced a refined formulation that became the standard in modern MoE architectures. This newer approach addresses some limitations of the importance-based loss by more directly measuring actual routing behavior rather than just router preferences. Instead of using coefficient of variation, it directly measures two quantities and penalizes their correlation.

The Switch Transformer formulation separates two components of balance: routing outcomes, meaning how many tokens end up at each expert, and router intentions, meaning how much probability mass the router assigns to each expert. A well-balanced system should keep both in sync. If routing is balanced but router intentions are skewed, the balance is fragile and may break down. If router intentions are balanced but routing is skewed, the discrete selection step behaves unexpectedly.

For a batch of $N$ tokens and $E$ experts, we define two complementary measures of expert usage. The first captures what happens during routing, while the second captures what the router intended. Their product forms the balancing loss.

Fraction of tokens routed to expert $i$:

$$f_i = \frac{1}{N} \textstyle \sum_{x \in \text{batch}} \mathbf{1}[\text{argmax}(g(x)) = i]$$

where:

• $f_i$: the fraction of tokens routed to expert $i$
• $N$: the total number of tokens in the batch
• $\mathbf{1}[\cdot]$: the indicator function, which is 1 if the condition is true and 0 otherwise
• $\text{argmax}(g(x))$: the index of the expert with the highest probability for token $x$
• $\textstyle \sum_{x \in \text{batch}} \mathbf{1}[\ldots]$: the total count of tokens in the batch assigned to expert $i$

This quantity measures the actual routing outcome. When the router makes its final decision for each token, selecting the expert with the highest probability, some experts will be chosen more frequently than others. The fraction $f_i$ counts how many tokens were ultimately sent to expert $i$ and normalizes by the total token count. This is the fraction of tokens where expert $i$ was selected (had the highest gate probability). It represents the actual load on each expert.

The important aspect of $f_i$ is that it reflects discrete decisions. A token is either routed to an expert or it isn't, and this binary nature means $f_i$ captures the ground truth of computational load. If expert 3 processes 40% of tokens, then $f_3 = 0.4$, regardless of whether those routing decisions were made with high confidence (probability 0.9) or narrow margins (probability 0.26 versus 0.25 for the runner-up).

Fraction of router probability assigned to expert $i$:

$$P_i = \frac{1}{N} \textstyle \sum_{x \in \text{batch}} g(x)_i$$

where:

• $P_i$: the fraction of router probability assigned to expert $i$
• $g(x)_i$: the routing probability assigned to expert $i$ for token $x$
• $N$: the total number of tokens in the batch
• $\textstyle \sum_{x \in \text{batch}} g(x)_i$: the sum of routing probabilities for expert $i$ across all tokens

This quantity measures the router's overall preference for each expert. Unlike $f_i$, which counts discrete decisions, $P_i$ aggregates the continuous probability values. This is the average probability the router assigns to expert $i$ across all tokens. It represents the router's "intention" to use each expert, capturing which expert wins the argmax competition and how confident the router is in its choices.

The distinction between $f_i$ and $P_i$ is subtle but important. Consider a batch where half the tokens give expert 1 probability 0.51 and expert 2 probability 0.49, while the other half gives expert 2 probability 0.51 and expert 1 probability 0.49. The fractions $P_1$ and $P_2$ would both be approximately 0.5. This reflects balanced intentions. But if by chance the argmax consistently favors expert 1, then $f_1$ would be much larger than $f_2$, revealing the imbalance in actual routing.

The load balancing loss is the scaled dot product of these vectors. Before writing the formula, note what we want to compute: a quantity that is small when both $f_i$ and $P_i$ are evenly distributed, and large when they concentrate on the same experts. The natural way to detect this joint concentration is to multiply the two vectors element-wise and sum:

$$\mathcal{L}_{\text{balance}} = E \cdot \sum_{i=1}^{E} f_i \cdot P_i$$

where:

• $\mathcal{L}_{\text{balance}}$: the auxiliary balancing loss
• $E$: the total number of experts
• $f_i$: the fraction of tokens routed to expert $i$ (actual load)
• $P_i$: the average probability assigned to expert $i$ (intended load)
• $\sum_{i=1}^{E}$: the summation over all $E$ experts

The factor $E$ normalizes the loss so that perfect balance gives $\mathcal{L}_{\text{balance}} = 1$. This normalization is a thoughtful design choice that makes the loss interpretable. A value of 1.0 means perfect balance, values above 1.0 indicate imbalance, and the magnitude tells us how severe the imbalance is. Under perfect balance, each expert receives fraction $1/E$ of both tokens and probability, so $f_i = P_i = 1/E$ for all $i$. We can verify the loss value:

$$\begin{aligned} \mathcal{L}_{\text{balance}} &= E \cdot \sum_{i=1}^{E} \left(\frac{1}{E} \cdot \frac{1}{E}\right) && \text{(substitute balanced values)} \\ &= E \cdot \sum_{i=1}^{E} \frac{1}{E^2} && \text{(simplify terms)} \\ &= E \cdot \left( E \cdot \frac{1}{E^2} \right) && \text{(sum constant $E$ times)} \\ &= 1 \end{aligned}$$

When load is imbalanced, both $f_i$ and $P_i$ grow for popular experts and shrink for unpopular ones. The correlation between these quantities drives the loss increase. Since both terms are larger for popular experts, their product is even larger, and the sum exceeds 1.

This loss formulation works because of what it penalizes. The product structure creates a natural feedback mechanism that targets exactly the behavior we want to discourage. Consider what happens when expert $i$ is overused:

1. More tokens are routed to it, so $f_i$ increases
2. The router assigns higher probabilities to it, so $P_i$ increases
3. The product $f_i \cdot P_i$ grows quadratically

This quadratic growth matters. When an expert is slightly overused, both $f_i$ and $P_i$ increase slightly, and their product reflects this twice. When an expert is heavily overused, both quantities increase substantially, and their product amplifies the signal further. The loss creates a gradient that pushes the router to reduce probabilities for overused experts.

Only $P_i$ contributes gradients, since $f_i$ involves a discrete argmax and is non-differentiable. But $f_i$ amplifies the gradient signal: popular experts get stronger pushback. The non-differentiable $f_i$ acts as a weighting factor, directing optimization toward high-load experts. Experts with high load receive stronger gradient signals on their probability terms, while experts with low load receive weaker signals. This adaptive weighting helps the optimizer prioritize the most severe imbalances.

The formulation also avoids a subtle failure mode. If we only penalized $P_i$ variance, the router could game the system by assigning nearly uniform probabilities while still routing most tokens to one expert through tiny differences that favor the same expert. Imagine a scenario where the router learns to give expert 1 probability 0.251 and all other experts probability 0.2497. The probability distribution looks nearly uniform, so a variance penalty would be small. Yet the argmax would consistently select expert 1, creating severe imbalance. By including $f_i$, we measure routing decisions rather than only intentions. The combined product $f_i \cdot P_i$ catches this gaming attempt because $f_1$ would be large even if $P_1$ increases only slightly.