Auxiliary Balancing Loss › Worked Example: Computing the Balancing Loss by Hand

To solidify the mechanics of the load balancing loss, let's work through a small concrete example step by step, using numbers that are easy to follow. We will use a batch of 8 tokens and 3 experts, and trace through the entire computation manually.

Imagine we have a batch of 8 tokens and the router produces the following probabilities for each token. Each row sums to 1.0 because the probabilities come from a softmax operation.

Notice that Expert 0 is selected most often (4 out of 8 tokens), while Experts 1 and 2 receive 3 and 1 tokens respectively. This is a moderately imbalanced distribution that our loss should detect.

Step 1: Compute $f_i$ (fraction of tokens routed to each expert)

Count the tokens assigned to each expert via the argmax, then divide by 8:

$$\begin{aligned} f_0 &= 4/8 = 0.500 \\ f_1 &= 3/8 = 0.375 \\ f_2 &= 1/8 = 0.125 \end{aligned}$$

Under perfect balance with 3 experts, we would expect each $f_i = 1/3 \approx 0.333$. Expert 0 is overloaded at 0.5, and Expert 2 is underloaded at 0.125.

Step 2: Compute $P_i$ (average router probability per expert)

Sum the probabilities in each expert's column, then divide by 8:

$$\begin{aligned} P_0 &= (0.70 + 0.60 + 0.15 + 0.10 + 0.55 + 0.20 + 0.65 + 0.25) / 8 \\ &= 3.20 / 8 = 0.400 \\ P_1 &= (0.20 + 0.25 + 0.65 + 0.70 + 0.30 + 0.15 + 0.25 + 0.60) / 8 \\ &= 3.10 / 8 = 0.3875 \\ P_2 &= (0.10 + 0.15 + 0.20 + 0.20 + 0.15 + 0.65 + 0.10 + 0.15) / 8 \\ &= 1.70 / 8 = 0.2125 \end{aligned}$$

We can verify: $P_0 + P_1 + P_2 = 0.400 + 0.3875 + 0.2125 = 1.0$. This must always hold because the softmax probabilities for each token sum to 1, so summing across all tokens and then averaging preserves this property.

Step 3: Compute the load balancing loss

We compute the dot product of $f$ and $P$, scaled by $E = 3$:

$$\begin{aligned} \mathcal{L}_{\text{balance}} &= E \cdot \sum_{i=0}^{2} f_i \cdot P_i \\ &= 3 \cdot (f_0 \cdot P_0 + f_1 \cdot P_1 + f_2 \cdot P_2) \\ &= 3 \cdot (0.500 \times 0.400 + 0.375 \times 0.3875 + 0.125 \times 0.2125) \\ &= 3 \cdot (0.2000 + 0.1453 + 0.0266) \\ &= 3 \times 0.3719 \\ &= 1.1156 \end{aligned}$$

The loss of 1.1156 is above the perfect-balance baseline of 1.0, confirming that the distribution is moderately imbalanced. The deviation is modest because the imbalance itself is moderate; if Expert 0 had received all 8 tokens, the loss would be much higher. Under perfect balance with $E = 3$ experts, we would have $f_i = P_i = 1/3$ for all $i$, and the loss would compute to exactly 1.0.

Intuition check: The product $f_i \cdot P_i$ is largest for Expert 0. It receives the most tokens and has the highest average probability. This is the "double penalty" mechanism: an overloaded expert contributes disproportionately to the sum because both factors increase.