Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

1. It is unclear if or how the gradients from auxiliary load balancing loss interfere with the model's training objective.

   We acknowledge that our assertion is primarily based on intuition and requires more rigorous experimental or theoretical validation. We will try to recheck this with more direct experiments and analysis in the next revision.

2. The argument about future token leakage in expert choice routing is unconvincing. The explanation lacks detail and Fig. 6 doesn't adequately demonstrate how this breaks causal constraints in language modeling. The empirical evidence assumes token leakage based solely on "abnormal loss drop."

   • Fig. 6 illustrates a potential information leak mechanism: When processing "cat" in position 2, the model can assign gate scores (0, 2) to make position 1 select expert 1 under Expert Choice. This creates a channel where position 1 can learn about the future token "cat" (e.g., by fixing expert 1 to output hidden states with high "cat" logits, so that when expert 1 selected, the output token must be "cat").

   • In practice, the mechanism may be more complex, potentially encoding future tokens as bits transmitted through multiple Expert Choice MoE layers. And we will emphasize in our revision that this is a potential risk rather than a guaranteed occurrence.

   • Regarding the 0.1 loss drop: While expert choice may offer genuine improvements, such a large gap is suspect. Our concerns align with findings in "Unified Scaling Laws for Routed Language Models" [1], which hypothesizes similar leakage risks.

     "Our hypothesis is that the re-assignment provides a subtle side channel through which information can be propagated backwards in time, and this can be abused by larger models resulting in the validation loss diverging during training. Adding a shuffling stage ameliorates this issue by introducing a large number of irrelevant elements to the rebalancing process, making it harder to infer behavior of future inputs."

     Appendix B.2.2 in [1]

   • Still, we agree with [1] that "further work is needed to confirm this theory."

3. The study's limitation to DeepSeekMoE instead of more widely used models like Mixtral or Switch transformer.

   Our focus on the DeepSeekMoE-like architecture can be well-justified for several reasons:

   • This design, combining fine-grained routed and shared experts, has demonstrated superior performance over traditional architectures (Mixtral/Switch transformer) in multiple studies: DeepSeekMoE [2], Qwen1.5 MoE [3], and OLMOE paper [4]. (You may find a detailed discussion in [4])
   • Apart from being successfully adopted by models like Qwen1.5 MoE and XVERSE-MoE-A36B [5], it has also proven highly effective in large-scale implementations such as DeepSeek-V2 (236B) [6] and DeepSeek-Coder-V2 (236B) [7]. Given this, we believe our experiments on this architecture provide sufficient evidence without needing to test old architectures like GShard / Switch Transformers / Mixtral.
   • Our implementation can be seen as an instantiation of modern MoE architectures, including DeepSeekMoE, QWen and OpenMoE [8], with architecture hyper-parameters optimized for our specific scale.
   • As mentioned below, we include experiments at larger scale (16B) / with more experts (256).

4. The performance gains are modest despite significant load balancing improvements.

   • Our improvements are statistically significant, exceeding typical random variation (<0.01). Figure 2 demonstrates that when comparing models with aligned $\text{MaxVio}_\text{global}$, we observe a substantial perplexity difference of approximately 0.2.

   • Our 16B-scale MoE model, trained for 1.33T tokens, shows steady improvements across several major benchmarks versus the baseline using auxiliary_loss_alpha=0.001:

     Benchmark	# shots	Our method	Vanilla Auxiliary Loss
     BBH	3-shot	39.31	37.31
     MMLU	5-shot	51.77	50.95
     C-Eval	5-shot	52.30	51.56
     CMMLU	5-shot	55.16	54.53

   • Additional experiments with 256 experts at 3B scale (increasing expert count while maintaining other 3B scale hyper-parameters) demonstrate that our method's advantages become more pronounced when the expert number increased:

     Method	Validation Perplexity (↓)	${**MaxVio**_**global**}$(↓)
     Vanilla Auxiliary Loss	7.38	1.13
     Our Method	7.23	0.10

1. The motivation and underlying intuition for the proposed approach could be clarified further to enhance understanding.

   We acknowledge this feedback and will enhance the explanation of Loss-Free Balancing's motivation and mechanics, particularly focusing on how it achieves an improved trade-off between load balance and performance.

2. Additional experiments are recommended to demonstrate the robustness of this approach when applied across varying numbers of expert mixtures. A scalable analysis would also be beneficial.

   We appreciate this suggestion and have conducted two additional sets of experiments:

   • At 3B scale, we increase the expert number to 256 (other hyperparameters unchanged), Our Loss-Free Balancing demonstrated enhanced advantage:

     	Validation Perplexity(↓)	${**MaxVio**_**global**}$(↓)
     Vanilla Auxiliary Loss	7.38	1.13
     Our Method	7.23	0.10

   • At 16B scale and 1.33T training tokens, we compare our method against the vanilla auxiliary loss baseline (auxiliary_loss_alpha=0.001), demonstrating steady improvements across multiple benchmarks:

     Benchmark (Acc.)	# shots	Our Method	Vanilla Auxiliary Loss
     BBH	3-shot	39.31	37.31
     MMLU	5-shot	51.77	50.95
     C-Eval	5-shot	52.30	51.56
     CMMLU	5-shot	55.16	54.53

3. The selection and initialization process in Formula 3.

   The process is currently stated in Algorithm 1 (page 4) with supporting explanations in the surrounding text. We will revise this section for greater clarity.

4. Regarding the stated improvement in load balance with increased computation batch size (Page 8): Is there an upper bound?

   As illustrated in Figure 5, our method and vanilla auxiliary loss exhibit different convergence behavior:

   • For our method, the MaxVio asymptotically approaches 0 as the batch size tends toward infinity. (the $\text{MaxVio}_\text{global}$ equals to a computation batch size = about 30k based on our testing protocol, for which our method is below 0.1)
   • For vanilla auxiliary loss, MaxVio may not come close to 0 when the computation batch increases. (the $\text{MaxVio}_\text{global}$ equals to a computation batch size = about 30k based on our testing protocol, for which the vanilla auxiliary loss still has high MaxVio.)

5. Is there a specific reason for selecting 9 MoE layers for testing? This hyperparameter was selected for our 1B baseline model and inherited for consistency in evaluating Loss-Free Balancing. (And for the 3B model, we use 11 layers)

6. Is any theoretical proof available for this MoE method? The inherent complexity of deep neural networks makes theoretical proofs challenging, not only for our approach but also for traditional vanilla auxiliary loss methods. In this paper, we prioritize experimental validation and have observed convincing improvements, as outlined both above and throughout our work.

1. Reference Scope

   We acknowledge that our focused discussion on MoE load balancing may have led to a limited reference scope. We will expand our literature review to include broader MoE-related works in the revised version. We welcome specific reference suggestions to enhance the comprehensiveness of our discussion.

2. Figure Adjustments

   We will optimize figure sizes in the revised manuscript.

3. Architectural Comparisons

   Our focus on the DeepSeekMoE-like architecture is well-justified for several reasons:

   • This architecture, featuring fine-grained routed experts and shared experts, proved to be superior by DeepSeekMoE [1], Qwen1.5 MoE [2], and OLMOE paper [3]. (You can find a detailed discussion in [3]). Apart from being successfully adopted by models like Qwen1.5 MoE and XVERSE-MoE-A36B [4], it has also proven highly effective in large-scale implementations such as DeepSeek-V2 (236B) [5] and DeepSeek-Coder-V2 (236B) [6]. Given this, we believe our experiments on this architecture provide sufficient evidence without needing to test old architectures like GShard [7] / Switch Transformers [8] / Mixtral.
   • Our implementation can be seen as an instantiation of modern MoE architectures, including DeepSeekMoE, QWen and OpenMoE, with architecture hyper-parameters optimized for our specific scale.
   • Regarding LLaMA-MoE: Recent studies (OLMOE [3], Skywork-MoE [9]) demonstrate that MoE models trained from scratch will ultimately outperform upcycled models after a decent number of training tokens, making it less relevant for our pretraining-focused study.

4. Performance Improvements

   The significance of our improvements is demonstrated through multiple lines of evidence:

   • The perplexity improvements exceed random variation (typically <0.01), indicating statistical significance. Figure 2 also shows a perplexity difference of about 0.2 when aligned with $\text{MaxVio}_\text{global}$.

   • Our 16B-scale MoE model, trained for 1.33T tokens, shows steady improvements across several major benchmarks versus the baseline using auxiliary_loss_alpha=0.001:

     Benchmark	# shots	Our method	Vanilla Auxiliary Loss
     BBH	3-shot	39.31	37.31
     MMLU	5-shot	51.77	50.95
     C-Eval	5-shot	52.30	51.56
     CMMLU	5-shot	55.16	54.53

   • Additional experiments with 256 experts at 3B scale (increasing expert count while maintaining other 3B scale hyper-parameters) demonstrate that our method's advantages become more pronounced when expert number increased:

     Method	Validation Perplexity (↓)	${**MaxVio**_**global**}$(↓)
     Vanilla Auxiliary Loss	7.38	1.13
     Our Method	7.23	0.10

[1] https://arxiv.org/pdf/2401.06066

[2] https://qwenlm.github.io/blog/qwen-moe/

[3] https://arxiv.org/pdf/2409.02060

[4] https://github.com/xverse-ai/XVERSE-MoE-A36B

[5] https://arxiv.org/abs/2405.04434

[6] https://arxiv.org/pdf/2406.11931

[7] https://arxiv.org/abs/2006.16668

[8] https://arxiv.org/abs/2101.03961

[9] https://arxiv.org/pdf/2406.06563

1. Is the improvement of Loss-Free Balancing significant? Could more benchmarks be reported?

   • The improvement is statistically significant, as the observed perplexity differences exceed typical random variation (which is usually below 0.01). Furthermore, Figure 2 demonstrates that when comparing models with aligned $\text{MaxVio}_\text{global}$, we observe a substantial perplexity difference of approximately 0.2.

   • We conducted evaluation on a 16B-scale MoE model trained with our method versus a baseline using auxiliary_loss_alpha=0.001. Both models were trained for 1.33T tokens, and our method shows steady improvement over the baseline (using auxiliary_loss_alpha=0.001) across multiple benchmarks:

     Benchmark (Acc.)	# shots	Our method	Vanilla Auxiliary Loss
     BBH	3-shot	39.31	37.31
     MMLU	5-shot	51.77	50.95
     C-Eval	5-shot	52.30	51.56
     CMMLU	5-shot	55.16	54.53

2. Can more models be reported?

   • Here, by the word 'DeepSeekMoE architecture', we refer to a mainstream MoE architecture design incorporating (a) fine-grained routed experts and (b) shared experts. This has been proved effective by 236B DeepSeek-V2 [1] and DeepSeek-Coder-V2 [2]. This architecture has also been adopted by current models including Qwen1.5 MoE [3] and XVERSE-MoE-A36B [4], and its benefits are thoroughly analyzed in the OLMOE paper [5] (you may refer to this for a detailed discussion). Given this, we believe our experiments on this architecture provide sufficient evidence without needing to test old architectures like GShard [6] / Switch Transformers [7].

   • Beyond the 16B-scale results presented above, we've conducted additional experiments with the expert number increased to 256 at 3B scale (increasing the expert number while keeping other hyperparameters the same as 3B setting), where Loss-Free Balancing enlarges its advantage on perplexity:

     	Validation Perplexity(↓)	${**MaxVio**_**global**}$(↓)
     Vanilla Auxiliary Loss	7.38	1.13
     Our Method	7.23	0.10

3. Is Max Violation a proper metric for MoE's load balance? Should auxiliary loss be considered as a load balance metric for Loss-Free Balancing?

   • Load balancing in MoE is primarily an engineering consideration rather than a mathematical abstraction. It's crucial for expert parallelization, where processing speed is bottlenecked by the most loaded device. Therefore, Max Violation directly measures what matters most in practical implementations.
   • Max Violation is more appropriate than auxiliary loss for measuring load imbalance because it directly reflects processing speed. Auxiliary loss, incorporating routing scores, doesn't purely measure load balance and lacks clear meaning (for example, two models with identical routing patterns may have different auxiliary loss values due to different routing scores).

4. For the "Questions" section: As 3. mentioned, the Max Violation metric effectively captures the load balance characteristics relevant to practical implementations. In the paper, we have reported our experimental results on the validation set (Section 4.2 & Section 5.1), which show the load balance situation of auxiliary-loss method and our method.

[1] https://arxiv.org/abs/2405.04434

[2] https://arxiv.org/abs/2406.11931

[3] https://qwenlm.github.io/blog/qwen-moe/

[4] https://github.com/xverse-ai/XVERSE-MoE-A36B

[5] https://arxiv.org/pdf/2409.02060

[6] https://arxiv.org/abs/2006.16668

[7] https://arxiv.org/abs/2101.03961