Auxiliary Balancing Loss › Why Gating Networks Cause Imbalance

Before diving into loss formulations, let's understand why the problem exists in the first place. The imbalance we observe in MoE models isn't a bug in the gating mechanism; rather, it's an emergent property of how neural networks learn through gradient descent. Understanding the root cause will make the solution much more intuitive. To see this clearly, consider what happens during early training when all experts and the router begin with random weights.

By chance, some experts receive slightly more tokens than others. This initial asymmetry might seem insignificant, perhaps one expert receives 27% of tokens while another receives 23%. However, this small difference sets in motion a cascade of effects. The experts that receive more tokens also receive more gradient updates, since each processed token generates gradients that flow back through the expert network. With more updates, these experts improve faster. Their weights adjust more quickly to the training distribution, and they become better at processing the token types they see.

The gating network, meanwhile, is doing its job diligently. It observes the reconstruction quality or task performance when tokens are routed to different experts, and it adjusts its weights to favor better-performing routes. When it notices that certain experts produce better outputs, it increases the routing probabilities toward those experts. This is exactly the behavior we want from a gating network when experts have different competencies, but during early training, this feedback creates a self-reinforcing cycle.

This creates a feedback loop: success breeds more success, failure breeds abandonment. Experts that received slightly more tokens initially improve, which makes them receive even more tokens and improve further. Meanwhile, experts that received fewer tokens fall behind. They update less frequently, improve more slowly, and become comparatively worse choices. The gating network, observing this growing gap, routes even fewer tokens to the struggling experts.

The gating network is doing exactly what we trained it to do: minimize task loss by routing tokens to the best experts. The problem is that "best" becomes self-fulfilling. Experts that receive no tokens never improve, so they remain poor choices, so they continue receiving no tokens. Left unchecked, this process converges to a degenerate state where one or two experts handle nearly all computation while the remaining experts contribute nothing. This represents a catastrophic waste of model capacity, as we've paid the memory cost of multiple expert networks but receive the computational benefit of only a few.

An auxiliary loss breaks this cycle by adding a cost for imbalance. The core insight is that we can modify what "success" means during training. Instead of optimizing purely for task performance, we add a secondary objective that penalizes concentration of expert usage. The total training objective becomes:

where:

• $\mathcal{L}_{\text{total}}$: the combined objective function used for training
• $\mathcal{L}_{\text{task}}$: the primary objective (e.g., cross-entropy loss)
• $\alpha$: a scalar coefficient controlling the weight of the balancing penalty
• $\mathcal{L}_{\text{balance}}$: the auxiliary loss term that penalizes imbalanced expert usage

The balancing loss increases when expert usage becomes skewed, forcing the optimizer to consider both task performance and expert utilization. This formulation allows us to tune how much we care about balance relative to task performance through the coefficient $\alpha$. A small $\alpha$ gently encourages balance while prioritizing task quality; a larger $\alpha$ enforces stricter balance at the potential cost of task performance.

The optimizer minimizes this combined objective. The gradient of the total loss with respect to any parameter now includes contributions from both the task loss and the balancing loss, weighted by $\alpha$. The router weights receive nudges from "how well did the selected expert perform" and from "is the distribution of expert selections becoming too skewed." These two gradient signals compete, and the coefficient $\alpha$ determines their relative influence.