Auxiliary Balancing Loss › Key Parameters

The key parameters for the auxiliary balancing loss implementation are closely interrelated, and understanding how they interact helps you make informed tuning decisions rather than treating each as independent.

• num_experts: Number of experts in the MoE layer. More experts increase the potential for imbalance (more chances for collapse) but also increase the model's capacity for specialization. As num_experts grows, you may need to be more careful with $\alpha$ tuning.
• aux_loss_coef: Coefficient $\alpha$ scaling the auxiliary loss (typically 0.01-0.1). This is the primary tuning knob. Start at 0.01 and adjust based on the imbalance ratio you observe during training.
• k: Number of experts selected per token (top-k routing). Larger $k$ naturally improves balance because each token contributes to more experts' loads. With $k=2$, the natural imbalance is typically lower than with $k=1$, which means you can often use a smaller $\alpha$.
• expert_size: Hidden dimension size of each expert network. Larger experts cost more compute per token, making load imbalance more expensive in practice. With larger experts, the economic argument for good balancing becomes stronger.