Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient

We thank the reviewer for their helpful feedback and comments. We also appreciate the recognition of the extensive empirical validation and actionable insights from our work. Below, we specifically address questions and weaknesses mentioned in the review. If our answers address the reviewer's concerns, we would like to kindly ask for the reconsideration of the rating.

Expert Imbalance To address the reviewer’s concern, we have performed an analysis of our training logs in the context of expert imbalance/token dropping. In general, we observe that the load balancing loss quickly induces balance between experts. Overall, the percentage of dropped tokens is low and doesn’t exceed 10%, therefore doesn’t significantly affect the training efficiency. Here (https://anonymous.4open.science/api/repo/3412679821-1326/file/dropped_toks.png?v=09990c3d) we present a plot of per-layer average amount of dropped tokens (excluding the first 10% of training), for 2 selected active parameter counts and number of experts varying from 2 to 32.

We will include an extended discussion of the token dropping in the final version of the manuscript. We thank the reviewer for this suggestion and believe this will be a valuable addition to our paper.

Fair Comparison Between Dense and MoE Models This is indeed a critical question. We took several steps when designing the experiments to ensure a fair comparison:

• During training, we use Switch MoE models with capacity factor C=1, i.e., tokens exceeding the capacity due to imbalance between experts are dropped. This is the most conservative setup for MoE; many papers [1, 2] and implementations [3, 4] actually use higher capacity factors or even dropless MoE. If we used such variants, we expect the benefits of MoE would be larger.
• We carefully adjust the learning rate depending on the number of experts and additionally tune the batch size.
• Finally, our work is the first to scale MoE models considering the total memory usage of the model, both during training and inference. This important comparison point was previously missing in the literature, which usually didn't consider the additional memory needed by MoE.

We hope that this explanation is convincing. Should the reviewer have further concerns or suggestions regarding this topic, we would be happy to address them.

Discussion on Scenarios With Limited Data Thank you for this important comment. Please note that since our analysis already considers different dataset sizes, our scaling law naturally applies to cases where the dataset size is fixed (by plugging the value of constrained dataset size $D$ in Eq. 6). Please refer to our response to reviewer njYw for a longer comment on this topic. We will also include a further discussion in the final version of the paper.

Routing Instability With the correctly tuned learning rate and batch size, as described in Sec. 5.1., we didn’t observe any instabilities caused by routing.

MoE Variants With Different Design Choices Thank you for the thoughtful question. Available literature shows robustness of scaling laws to changes in architecture/setup. This is documented for routing algorithms [6], training datasets [7] or depth-to-width ratio [5]. Therefore, while we agree that variations on the MoE design - such as the routing policy - form an important axis of model architecture, we expect our conclusions to be robust to these changes.

Fine-Grained/Coarse-Grained Experts We didn’t explicitly model expert granularity, since we would need vastly more resources to consider another variable in our experiment grid. However, based on the available literature [9, 10], we can expect that using fine-grained experts would further improve the efficiency gains we observe when using MoE. We leave further quantification of these gains for future work. We also hope to make it easy for the community to extend our research, since we will release the model checkpoints and code upon the end of the review period.

Regarding Other Comments In the paper, we always use the notation N_act to refer to active parameters, and N_total to refer to total parameters.

References:

[1] Muennighoff et al., OLMoE: Open Mixture-of-Experts Language Models

[2] Vavre et al., Llama 3 Meets MoE: Efficient Upcycling

[3] Gale et al., MegaBlocks: Efficient Sparse Training with Mixture-of-Experts

[4] Tan et al., Scattered Mixture-of-Experts Implementation

[5] Kaplan et al., Scaling Laws for Neural Language Models

[6] Clark et al., Unified Scaling Laws for Routed Language Models

[7] Hoffmann et al., Training Compute-Optimal Large Language Models

[8] Frantar et al. Scaling Laws for Sparsely-Connected Foundation Models

[9] Ludziejewski et al., Scaling Laws for Fine-Grained MoE

[10] Dai et al., DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models

We thank the reviewer for the insightful feedback and comments. We appreciate the recognition of the vast scale of our experiments, clear visualizations, and good communication of findings. Below we address the reviewer’s questions in details. If our answers address the reviewer's concerns, we would like to kindly ask for the reconsideration of the rating.

Variable Dataset Size Thank you for this valuable insight. Please note that in the scaling laws literature, it is common to treat the dataset size as a variable [1,2,3,4]. These results can be crucial in the long run, as shown by [2], which shifted the community’s focus to collecting larger datasets rather than just scaling model sizes. Furthermore, typically a one-epoch setting is assumed, i.e. the models are trained on the whole dataset. However, we agree that the data-constrained scenario is important from the practical perspective, and including such discussion can provide a broader context to our paper. Since our analysis already considers different dataset sizes, our scaling law naturally applies to cases where the dataset size is fixed (e.g., in Fig. 2 (a) in the paper each vertical line represents a constant dataset size). We will add a more detailed section concerning the dataset constraints in the camera-ready version of the paper.

Different Datasets [2] note that their scaling laws are robust to changes in the dataset (they consider 3 different datasets). Similarly, we expect that our qualitative conclusions will hold for different datasets, even though the numerical coefficients will likely be different.

Reliance on Chinchilla / Changes of Chinchilla for MoE As Chinchilla is an established result and is proven to reliably model scaling in various scenarios [3,4], our scaling law needed to be reducible to Chinchilla in the case of $E=1$ (dense). Our experiments have shown that our scaling law works well for different $E$’s via low MAE on a held-out extrapolation validation set, both for dense and MoE models. Note that we do modify Chinchilla scaling laws - a scaling law that generalizes [2] for a given E is a major contribution of our work. As a further piece of evidence for the reliability of our approach we will include a bootstrapped version of our results in the final version of the paper.

Expert Imbalance Thank you for raising this important question. To avoid exceeding the response length we refer to the answer given to reviewer V5XV.

Influence of Architecture/Training Optimizations Introducing changes (such as MLA) to the models’ architecture would change the scaling laws coefficients, however, we expect the functional form to remain same (similarly to [5], whose laws’ form did not change when they modified their routing mechanism). As changes to the attention mechanism would impact both Dense and MoE models, we would not expect them to favor any of the architectures, and we would expect the general conclusions to stay the same. Simultaneously, Switch MoE layer could be improved e.g. via the use of fine-grained experts, yielding better scaling behavior in MoE models.

FLOPs and Efficiency To compare our FLOPs estimates with actual numbers, we perform an experiment for 3 model sizes and $E$’s using torch.utils.flop_counter.FlopCounterMode, measuring MFLOPs per token in forward pass. The discrepancy between our estimate and actual numbers stems mostly from the implementation details of the embedding layer and becomes relatively smaller with larger models. It is also the same size in MoE and dense models with the same number of active parameters. Note as well that our FLOPs estimation method is standard in literature [2,6].

In a specific setup, MoE efficiency depends on the implementation, with many efficient ones [7,8] speeding up training and inference.

Distributed Training Since our method is implementation agnostic, distributed setting should not have any impact on our conclusions.

Other Comments Thank you for your comments regarding equation number and clarification on the FLOPs counting. We will fix them in the final version of the paper.

References:

[1] Kaplan et al., Scaling Laws for Neural Language Models

[2] Hoffmann et al., Training Compute-Optimal Large Language Models

[3] Ludziejewski et al., Scaling Laws for Fine-Grained Mixture of Experts

[4] Kumar et al., Scaling Laws for Precision

[5] Clark et al., Unified Scaling Laws for routed language models

[6] Gadre et al. Language models scale reliably with over-training and on downstream tasks

[7] Tan et al., Scattered Mixture-of-Experts Implementation

[8] Zhao et al., DeepEP: an efficient expert-parallel communication library

We thank the reviewer for the detailed comments and suggestions. We also appreciate the recognition of the practical importance of our findings, the actionable insights they provide, and the confirmation that our claims are well-supported. We hope that the answers below adequately answered the reviewer's questions and concerns. If that is the case, we kindly ask for a reconsideration of the paper score.

Regarding the innovations, differences and challenges compared to previous work We believe that our paper delivers significant innovations not present in existing works and provides important value to the scientific community:

• Crucially, we derive and fit a new scaling law, which allows us to develop novel, practical insights for MoE training. In particular, we are the first to consider scaling of MoE models in a memory-constrained regime. We deliver a new, unexpected result that MoE models can be optimal in such scenario. We are the first to obtain the optimal token-to-param ratio based on the number of experts.
• The inclusion of D allows us to reach qualitatively different conclusions than [5]. Furthermore, we consider the compatibility of our formula with [6] a strong point of our paper.
• We have performed a careful and meticulous work of collecting empirical evidence (>280 training runs, spanning model sizes with up to 5B parameters, multiple distinct training durations and MoE expert counts). These results set a substantial basis for deriving conclusions. Moreover, we open the results to the scientific community for further analysis, by releasing model checkpoints and training code upon the end of the review period.
• Furthermore, we put a lot of attention towards details to ensure our results are robust. We utilize a dynamic batch size (Sec. 5.1.1.) to ensure the reliability of our training runs regardless of their token count. Additionally, we derive a joint scaling law for the learning rate (Sec. 5.1.2, App. C), which is a novel contribution of our work. Based on the fitted coefficients, it justifies that a larger $E$ necessitates lower learning rate - a result not present in the literature to this point.

Evaluation As we are interested in comparing general trends, we focus on modeling perplexity. This metric has been shown to predict downstream performance well, even if the architecture details (e.g. the model size) differ [1]. Existing literature [2], suggests that when perplexity is fixed, MoE outperforms dense models at world-knowledge tasks while matching their performance in reasoning. Notwithstanding, we agree that an analysis of the downstream performance would be a valuable addition. We have performed such experiments (Results: https://anonymous.4open.science/api/repo/3412679821-1326/file/all_benchmarks_grid.png?v=70a33096) using dropless MoE during evaluation. We find that the perplexity strongly dictates the overall downstream performance, however there seems to be a slight advantage of either dense or MoE models in selected benchmarks (LAMBADA, OpenBookQA). We will add the analysis of downstream performance in the camera-ready version of the paper.

Derivation of the Optimal N and D Here (https://anonymous.4open.science/api/repo/3412679821-1326/file/optimal_n_d.png?v=331aaf30) we present a sketch of the derivation of the optimal N_act (with optimal D being analoguous). We will provide full details in the final version of the manuscript. We thank the reviewer for this suggestion and believe it will contribute to the completeness of the analysis.

Dataset Selection We train our models using FineWeb-Edu, a 1.3T subset of the FineWeb dataset - a large, openly available, high-quality LLM pretraining dataset. The data curation process of the FineWeb was guided using popular benchmarks (CommonSenseQA, HellaSwag, etc.). FineWeb-Edu is selected using a filter for highly educational content. It “outperforms all openly accessible web-datasets on a number of educational benchmarks” [3]. We will clarify our choice of the dataset in the camera-ready version.

Regarding “Questions for the Authors” Although we focus on the standard MoE variant, we believe that our main conclusions will hold for other MoE versions. We can form this assumption based on related work, where scaling laws are shown to be consistent across routing algorithms [5] or datasets [6]. Based on the literature [4], we can expect changes like fine-grained experts to further improve efficiency gains from using MoE.

References:

[1] Du et al., Understanding Emergent Abilities of Language Models from the Loss Perspective

[2] Jelassi et al., Mixture of Parrots: Experts improve memorization more than reasoning

[3] Penedo et al., The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale

[4] Ludziejewski et al., Scaling Laws for Fine-Grained Mixture of Experts

[5] Clark et al., Unified Scaling Laws for routed language models

[6] Hoffmann et al., Training Compute-Optimal Large Language Models

Thank you for your thoughtful and encouraging review. We are especially grateful for the reviewer’s recognition that "the key contributions of the paper are of high importance to the community of language modeling" and that the work is of "great quality". We are also thankful for the detailed comments and suggestions, which we will apply in the revised manuscript version.

Scaling Laws for Other Modalities Thank you for pointing this out. We agree that the techniques presented in our work could be extended to other domains such as computer vision, and we appreciate the references to Riquelme et al. (2021) and Zhai et al. (2021). We will incorporate these works in the revised version and expand the discussion to highlight the broader applicability of our scaling laws beyond language modeling.

Regarding "Other Comments" If we are able to fit the labels within Table 2, we will do so in the camera-ready version.

Regarding the typo in Fig. 5 - you are absolutely right, the caption should indicate that D/N < 1 stands for fewer tokens than parameters and we will fix the caption.

Thank you for pointing out the missing label in Fig. 5b. The brown curve indeed corresponds to 16B training tokens - we will fix the plot accordingly.