Auxiliary Balancing Loss › Loss Coefficient Tuning

The coefficient $\alpha$ in $\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \alpha \cdot \mathcal{L}_{\text{balance}}$ controls the strength of the balancing incentive. This single hyperparameter has outsized importance in MoE training, as it determines how the model allocates its optimization effort between learning useful representations and maintaining computational balance. Setting $\alpha$ correctly is important: too small and experts collapse; too large and the model sacrifices task performance to achieve perfect balance.

To understand why $\alpha$ matters so much, think about what gradient descent is doing at each update step. The gradient of $\mathcal{L}_{\text{total}}$ has two components: one from the task loss that pushes the router toward routing decisions that improve task performance, and one from the balancing loss that pushes the router toward more uniform distributions. The coefficient $\alpha$ sets the ratio of these two forces. If $\alpha = 0.01$ and the task loss gradient magnitude is 100 times larger than the balancing loss gradient magnitude, then both contributions are comparable in their effect. If $\alpha$ is too small, the balancing force is drowned out and has negligible effect on the router's learned behavior.

Typical values used in practice and the reasoning behind them:

• $\alpha = 0.01$: The default value from the Switch Transformer paper. Works well for most MoE configurations. At this setting, with a typical task loss around 2.0 and a balancing loss around 1.0, the balancing term contributes roughly 0.5% of the total gradient signal.
• $\alpha = 0.001$: Appropriate for models where task quality is necessary and some imbalance is acceptable, or as a starting point for exploratory training.
• $\alpha = 0.1$: More aggressive balancing. Use this if you observe persistent collapse with smaller values.
• Scale dependence: The appropriate $\alpha$ depends on how large $\mathcal{L}_{\text{task}}$ typically is. If task loss is around 2.0 and balancing loss is around 1.0, $\alpha = 0.01$ means balancing contributes about 0.5% of the total gradient signal.
• Number of experts: More experts require careful tuning, as the minimum achievable imbalance increases with expert count.
• Top-k value: With top-2 routing, each token contributes to two experts' load, which naturally improves balance compared to top-1.

When $\alpha$ is too small, the model effectively ignores the balancing signal. The gradient contribution from the auxiliary loss becomes negligible compared to the task loss gradients. Expert collapse proceeds as if no auxiliary loss existed. You'll observe a few experts receiving most tokens while others remain unused throughout training. The training curves may look healthy in terms of task loss, but inspection of expert utilization reveals the underlying pathology. Essentially, the optimizer "prices in" the tiny balancing penalty and determines that the task loss improvement from putting tokens is worth it. You've told the model that balance matters a little, but the model has correctly concluded that it doesn't matter enough to change its behavior.

When $\alpha$ is too large, the model prioritizes balance over task performance. The optimizer, faced with a strong gradient signal to equalize expert usage, adjusts the router weights primarily to reduce the balancing loss. In the extreme, the router learns to assign exactly uniform probabilities to all experts, achieving perfect balance but terrible task performance. The model loses its ability to specialize experts for different types of inputs. Every expert becomes equally mediocre at handling all tokens, squandering the capacity benefits that MoE architectures are designed to provide. In a language model, this would look like a model that routes all token types to all experts equally, preventing any expert from developing useful specialization for code, dialogue, or technical text.

The sweet spot achieves "soft" balance: experts receive roughly similar token counts, but the router retains freedom to assign non-uniform probabilities based on input characteristics. In this regime, the auxiliary loss prevents pathological collapse without overly constraining the router's flexibility. You can usually identify this regime by checking that the imbalance ratio (max expert usage divided by min expert usage) stays below about 3:1, while the router's probability distributions for individual tokens remain diverse rather than uniform.

Tuning $\alpha$ requires monitoring both the task loss and expert utilization during training. You cannot rely on validation loss alone, because a model with severe expert collapse can sometimes achieve decent task loss by putting all computation in one or two strong experts. The hidden cost appears when you measure inference efficiency or examine whether the model uses its full capacity. A reliable approach is to start with $\alpha = 0.01$ and monitor both metrics during training:

1. Track expert utilization (fraction of tokens per expert) across training
2. Track the ratio $\mathcal{L}_{\text{balance}} / \mathcal{L}_{\text{task}}$
3. If experts collapse (one expert gets >50% of tokens), increase $\alpha$ by 2-3x
4. If task loss plateaus while balance loss keeps dropping, decrease $\alpha$

We'll explore a related technique called Router Z-Loss in the next chapter, which provides additional stabilization with less impact on task performance.