Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient
Authors: Jan Ludziejewski
Maciej Pióro
Jakub Krajewski
Maciej Stefaniak
Michał Krutul
Jan Małaśnicki
Marek Cygan
Piotr Sankowski
Kamil Adamczewski
Piotr Miłoś
Sebastian Jaszczur
Venue: ICML 2025 poster (2025)
Keywords: mixture of experts
scaling laws
llm
Primary area: deep_learning->large_language_models
TL;DR: We derive scaling laws for memory- and compute-constrained Mixture of Experts language models.

Abstract: Mixture of Experts (MoE) architectures have significantly increased computational efficiency in both research and real-world applications of large-scale machine learning models. However, their scalability and efficiency under memory constraints remain relatively underexplored. In this work, we present joint scaling laws for dense and MoE models, incorporating key factors such as the number of active parameters, dataset size, and the number of experts. Our findings provide a principled framework for selecting the optimal MoE configuration under fixed memory and compute budgets. Surprisingly, we show that MoE models can be more memory-efficient than dense models, contradicting conventional wisdom. Extensive empirical validation confirms the theoretical predictions of our scaling laws. These results offer actionable insights for designing and deploying MoE models in practical large-scale training scenarios.

Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient
Jan Ludziejewski" 12 12 Maciej Pióro 23 Jakub Krajewski "
Maciej Stefaniak
Michal Krutul 12 Jan Malaśnicki 12
Marek Cygan 14 Piotr Sankowski 516 Kamil Adamczewski
Piotr Miloš 28
Sebastian Jaszezur
Abstract
widely adopted (Jiang al., 2024; Dai al., 2024) and are generally considered compute-efficies (Ludziejewski et al., 2024; Clark et al., 2022), though often considered to be infffiiiiee (Zadourí et al., 2024). However, the precise trade-offs between compute and memory efficiency remain unclear.
Mixture of Experts (MoE) architectures have significantly increased computational efficiency in both research and real-world applications of large-scale machine learning models. However, their scalability and efficiency under memory constraints remain relatively underexplored. In this work, we present joint scaling laws for dense and MoE models, incorporating key factors such as the number of active parameters, dataset size, and the number of experts. Our findings provide principled framework for selecting the optimal MoE configuration under fixed memory and compute budgets. Surprisingly, show that MoE models can be more memory-efficient than dense models, contradicting conventional wisdom. Extensive empirical validation confirms the theoretical predictions of our scaling laws. These results offer actionable insights for designing and deploying MoE models in practical large-scale training scenarios.
Consider motivating question: Is an MoE model the op- timal choice when constrained by fixed memory budget, such single H100 node? While computational efficiency is important, it does not directly determine the optimal num- ber of experts. Increasing the number of experts has minimal impact on computation but can drastically raise memory re- quirements, often to a prohibitive level.
To address this question, we derive a joint scaling law for both dense and MoE models, accounting for key factors such as the number of active parameters, dataset size, and number of experts. This framework provides rigorous analysis of model performance under strict memory constraints. Our findings reveal that, contrary to common assumptions, MoE models can be more memory-efficient than dense models.
Our work is the first to provide detailed guidance se- lecting the optimal number of experts for MoE models, balancing both computational and memory constraints, Our conclusions are based on extensive, large-scale experiments comprising 270 models, scaled up to 5B parameters."
1. Introduction
Recently, language models have grown increasingly large, a trend accelerated by Mixture of Experts (MoE) tech- niques (Fedus et al., 2022; Du et al., 2022). MoE models are
In summary, the key contributions of this work are:
We derive a joint scaling law for Mixture of Experts and dense models.
"Core contributors University of Warsaw VIDEAS NCBR IInstitute of Fundamental Technological Research, Polish Academy of Sciences *Nomagic 'Rescarch Institute IDEAS "MIM Solutions 'Wroclaw University Seience and Technol- ogy "Institute of Mathematics. Polish Academy of Sciences. Correspondence to: Jan Ludziejewski
C(Nas. D) a Na+ssu(E)
+ + Det + (I) where L is the final training loss, Nact is the number of active parameters, D is the datasct size, E is the monotonic transformation of the number of experts, and is the irreducible entropy of the dataset.
Maciej Pióro <maciej, waaiii@mmaal.........m Jakub Krajew- ski egim.jakuble gomi Sebastian Jaszezur <sebas- tian.jaszezur @ gmail.com>.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
'Checppoints and inference code are available Hugging Face. Codebase used the experiments be found on GitHub.

---

Joint Scaling Laws for Mixture of Experts
2.35
Experts
2.75
Experts (dense)
2.3
2.7
E=1 Compute Optmal
2.25
16
Loss
Loss
2.2
2,00
2.15
2.1
10B Total Parameters
100B
2,56
Ax10" 6x101 8x10
1x20" 12x10
FLOPs
(a)
(b)
Figure 1. (a) The loss of memoryy-cnnstriined models predicted using scaling law under fixed training budget of 10 FI.OPs. Each curve represents different number of experts. The lines are truncated compute-optimal points, as undertrained models are both bigger and worse in terms of loss, thus pointless in -onstrrc-on scenario. Shaded areas present memory optimal number of experts for the corresponding parameter budgets. (b) Experimental validation the thesis that MoE can be memory optimal. The marked area shows interval in which training compute-matcher matched MoE achieves better loss than overtrained dense model with the same number of total parameters (1.1B). The resulting MoE was trained for longer and had less active parameters, making it more practical.
Based on the proposed scaling law, we show that the choice of the optimal number of experts (including dense models with E = 1) depends on specific com- putational and memory constraints, see Figure 1. Fur- thermore, demonstrate how the optimal token-to- parameter ratio depends on E.
by the development of extremely large-scale models such as DeepSeck-RR R1 and Qwen2.5-Max DeeepSeck-Al et al., 2025; Team, 2024a). Apart from language, MoE has also been shown to be an effective measure of scaling vision (Riquelme al., 2021).
In work, et al., 2022), which routes each token to one expert and encourages even token-to-expert assignment via the addition of a differentiable load-balancing loss.
the standard Switch MoE layer (Fedus
We show that MoE can often be the preferred alter- native to dense models, even GPU memory is the constraining factor. We validate our theoretical find- ings by training set of I.IB I. parameter models under identical compute and total memory budgets. The MoE models achieve lower final loss, confirming their su- perior efficiency in practice. Furthermore, we observe that MoE models only have lower loss, but deliver higher performance during inference.
Scaling Laws. Scaling laws refer to empirically derived equations that relate model loss to factors such the number of parameters, the quantity of training data, or the computa- tional budget. For dense Transformers, scaling laws were initially explored by Hestness et al. (2017) and Kaplan et al. (2020), who identified power-law relationships between the final loss, model size, and dataset size. Hoffmann et al. (2022) expanded this by incorporating variable cosine cycle lengths and adjusting the functional form of the equation:
2. Related Work
Mixture of Experts. Mixture of Experts (MoE) was in- troduced by Jacobs et al. (1991), who combined gat- ing network with a set of expert networks. Shazeer et al. (2017) applied MoE to an LSTM-based model (Hochreiter Schmidhuber, 1997), scaling the architecture up 137 billion parameters. In Transformer-based LLMs, MoE is most often applied replacement for the feed-forward layer (Lepikhin et al., 2020; Shazeer et al., 2018). 8). It re- places the feed-forward's MLP with set of expert MLPs along with a router, which selects one or more MLPs for each token. With the recent surge in LLM research, MoE models are gaining even more traction. This is exemplified
L(N D) = mN + nD +
Scaling laws have also been applied to other architectures and training setups. Zhai et (2022) derived scaling laws for Vision Transformers, scaling the models up to 2B pa- rameters. Henighan et al. (2020) examined autoregressive modeling across multiple modalities, while Ghorbani et al. (2021) focused on machine translation. Frantar et al. (2023) studied the effects of pruning on vision and language Trans- formers, determining optimal sparsity given a fixed compute budget.
2

---

Joint Scaling Laws for Mixture of Experts
2.7
Experts
Experts
2
2.6
Savings from switching to MoE
60%
8
16
2.5
50%
2.4
FLOPs Budget (Optimal Tokens)
40%
2.3
* 1e+20
30%
Se+20
20%
108
1008
1e18
1e19
1e20
1e21
1e22
1e23
Training Tokens
Dense Model FLOPs
(a)
(b)
Figure 2. (a) IsoFLOP profiles for selected training budgets. Compute optimal points marked for each (b) Savings from switching from compute-optimal dense model to MoE. The advantage of using MoE increases with larger models and expert counts.
Clark et al. (2022) investigated scaling in MoE models, varying model size and the number of experts on a fixed dataset, and concluded that routed models are more efficient only up to a certain size. Their formula took the form:
et al., 2023; Kumar et al., 2024; Ludziejewski et al., 2024), we postulate the following form of the equation:
C(N D) = m(E) n(E)D"(F) + d(E), (5)
((Nact E) = Not
No
(3)
assuming that if we fix the number of experts the model performance can be described using Equation 2. In the subsequent part, we will postulate how m, μ,, n, v, c depend on E, deriving the joint equation.
where É is monotonic transformation of the number of experts E defined as:
Constant Factor. c(E) represents irreducible loss caused by the inherent entropy of the dataset. Thus, it does not depend on the architecture (E in our case):
1
(4)
Fmaxx
E-1+
Exn
Eus
These analyses have since been extended by Ludziejewski al. (2024) and Dai et al. (2024), who considered variable dataset size as well as the granularity of experts. In our work, we keep the experts non-granular; however, we treat the number of experts and the number of training tokens as variables. Sardana et al. (2024) assumes a fixed joint infer- ence and training budget. We make similar assumptions; however, consider accelerator memory tor and extend the analysis to MoE models, which serve as more compute-friendly alternative to dense models. Yun et al. (2024) have focused on MoE inference optimality and measuring real hardware efficiency.
E
Interaction of E with Model and Dataset Size. To quan- tify the interaction between the number of experts and other training parameters, gather observations from related work:
limiting fac-
1. Scaling in E can be described as a power law (Clark et al., 2022).
2. For fixed dataset size, as model size increases, the benefit of using an MoE diminishes (Clark et al., 2022).
For a fixed model size, the number of training tokens increases, the benefit of an MoE grows (Ludziejewski al., 2024).
3. Joint MoE Scaling Laws
We now derive the functional form of our joint scaling laws for both dense Transformers and MoE, relating the number of active model parameters Nact: training tokens D, and MoE experts E.
Motivated by Observation 1, we set
m(E) o. n n(E) = bE".
Fixed Number of Experts. Following Hoffmann et al. (2022) and established practice in the literature (Frantar
reflecting the power-law relation between and the loss.
3

---

Joint Scaling Laws for Mixture of Experts
Additionally, to ensure flexibility in modeling Observations 2 and 3, we introduce interaction with the exponents over Nact and D:
Table I. Example compute-optimal training configurations for MoE models. For every training budget as the number of experts increases, the optimal D also goes up while Na decreases.
(( E) = a + (I((E) v = 3 Çln(E.
Training Budget
Experts
Napt act
Dope
1 x 1020
1
1.7B
9.7B
Note that if we ignore the second and third terms in Equa- tion 5, this yields functional form identical to Equation 3. Empirically, we observe a good fit for our formula, as de- scribed in Section 5. This shows that our proposed inter- actions between E, Nac, and D can accurately model the performance of MoE models.
2
1.5B
11.4B
4
1.2B
13.9B
8
990M
17B
16
810M
20.7B
5 x 10
1
4B
21B
2
3.5B
24B
4
3B
28B
Modeling of E. When the number of experts is small, certain overhead, caused, for example, by interference from auxiliary losses, can overshadow the benefits of conditional computation. Additionally, using very large numbers of experts brings diminishing returns. To account for these phenomena, we follow Clark et al. (2022) and use trans- formation of the number of experts É given in Equation 4.
8
2.5B
33.2B
16
2.1B
39B
1 x 1021
1
5.7B
29.3B
2
5B
33B
4
4.4B
38B
8
3.8B
44.3B
16
3.3B
51.2B
Joint MoE. Scaling Law. By combining these observations, we establish the final form of our scaling law:
C(N D, E) = aE N (artll())
+ +b D+t + c. (6)
model size. IsoFL.PP slices comparing the predicted loss with dataset size for selected compute budgets are plotted in Figure 2 (a).
We fit the coefficients in Equation 6 based on the results of experiments; Table 3. In Section 4, we present the outcomes and findings derived from the scaling laws. The details of the training runs, as well as the fitting procedure, are described in Section 5.
For any fixed F our scaling law has the Chinchilla functional form of Equation 2. Thus, from Hoffmann et al. (2022), the compute- crapute-optimall number of tokens and active parameters for the budget F and the number of experts E are given by
FF.
F (
4. Compute and Memory Optimality
NoR (F) G
D G-1
(9)
6
In this section, we employ our scaling laws to derive recommendations optimal settings in various training and inference scenarios. See Appendix A for details on counting FLOOPs, the relations between active and total parameters, and other technical details.
where
G ((F)m(()
((E)n(E)
and
4.1. Compute Optimality
v E p(E) + v(E)'
b M(E) ((E) + +v(E)
A model is considered compute-optimal if, among models trained with the same compute budget F, it achieves the lowest loss. To find such an optimal configuration, we
optimize the following:
The full derivation of Nopi Dep can be found in App.C. We compare the optimal configurations for several compute budgets in Table 1.
arg min C(Nas, D, E) D,E
(7)
s.t. 6. =
Both from comparing the IsoFLOP slices (Figure 2) and the values listed in the table, we can see that the compute- optimal configuration for given compute budget clearly depends on F, with MoE models requiring comparatively larger datasets and correspondingly smaller numbers of ac- tive parameters.
Optimal N and D Depend on the Number of Experts. Assuming a given number of experts E, the compute- optimal training configuration can be achieved by select- ing the appropriate trade-off between training tokens and

---

Joint Scaling Laws for Mixture of Experts
(Figure 1). Significantly, the MoE models attains lower loss even if the dense model is overtrained (i.e., after passing its compute-optimal token count).
Finding 1.
More experts - higher tokens-to-param ratio. Assume fixed compute budget. In this scenario, when increasing the number of experts, it is optimal de- crease the number of active parameters and increase the number of training tokens accordingly (Table I).
Finding 3. MoE also be memory-efficient. A notal-parameter-mattiee MoE model can outper- form dense model trained with the same compute budget (Figure 1). Moreover, such an MoE model is more compute- and memory-efficiet at inference.
Mixture of Experts is Compute Optimal. Now, we com- pare the performance across various numbers of experts, with respective values of tokens and active parameters op- timized. As illustrated in Figure 2, we observe significant compute savings for MoE models compared to dense mod- cls, with larger number of experts providing more pro- nounced benefits.
4.3. Total Memory Optimality
During autoregressive generation, a decoder-only model processes single token while storing activations (keys and values) for previous tokens in the KV cache. In the case of multi-head attention, its size equals 2T x NNooocks X dmodal- where T is the number of tokens in the cache (possibly within multiple sequences in the batch). Including the cache size yields the optimization criterion:
Finding 2. More experts >> better performance. For a given compute budget, increasing the number of experts always improves performance, provided the size of the model and the number of training tokens are adjusted (Figure (a)).
arg min L(Naa, D, E)) D,.DEE s.t. 6.N D= Nivia + 2T motocks S M
The higher efficiency of MoE in terms of training compute comes at price of increased memory requirements. How- ever, somewhat surprisingly, we find that MoE models can outperform dense models the same size trained with the same amount of training compute- a result we describe in more detail in the next subsection.
For practical values of T, fair comparison of memory requirements should include the size of KV cache in addition to the model size. Figure (b) presents the optimal models for a given compute and varying memory constraints when the size of the KV cache included. Importantly, MoE models compare more favorably dense models in this graph, and as T increases, they outperform dense models at ever smaller model sizes. In Figure (b), the E {2,4] models employ smaller KV cache. It means that if memory is constrained, the MoE model can store longer contexts or work with larger batch size than the dense model.
4.2. Model Memory Optimality
Often, is insufficient to consider models solely from the perspective of compute optimality, a compute-optimal model can be impractically large, preventing its deployment on available hardware. Additionally, a model can be in- efficient when run on a GPU with a small batch size (He, 2022). It thus natural to consider a straightforward exten- sion to the notion of compute optimality, specifically model memory optimality. We say a model memory optimal if, among models trained with the same compute budget F and having at most M parameters, it achieves the lowest loss:
4.4. Inference Optimality
Large models, while capable, might also be too costly to run due to their high computational demand. To account for this drawback, we can further assume that a model will process some number of tokens, Dints throughout its lifetime and find the best model whose demands do not exceed some predefined joint training and inference budget:
arg min L( E) D.E
S.t. 6N F. Nital S M
arg min L(Naat D, E) N.. s.t. 6. + 2N Dant F.
Note that model memory- differ in the number of active parameters- fraction of them. Intuitively, it should thus have worse performance. At the same time, given some budget, it can be trained on more tokens, lowering the loss. Our scaling laws suggest that MoE models can be model memory efficient. We validate this claim by training a 1.1B dense model and a model size and FLOP matched = (2, 4} counterparts
dense and MoE models
uses just
Figure 3 (c) presents the optimal models for given compute and varying memory constraints if a joint budget needs to accommodate both training and inference demands. We find that in this scenario, MoE models outperform dense
5

---

Joint Scaling Laws for Mixture of Experts
Experts
Experts
Experts
212
2.12
211
16
2.05
2.0
2.00
2.06
Tataal Parameters
Total Parameters
Total Parameters
(a)
(b)
(c)
Figure 3. Loss predicted for various expansion rates at FLOPs budget F = 5c22. The x-axis denotes the size of the corresponding dense model, possibly with KV cache. (a) The model size is simply the number of parameters. (b) The model size includes the KV cache (assuming 8192 tokens), (c) Additionally to KV cache, the training budget reduced by the inference cost on 100B tokens.
2.4
32
2.4
2.4
18
2.3
9
16
Logg
2.3
16
Loss
16
9
2.3
8
2.2
2.2
28
2.2
28
98
58
10B
2.1
2.1
2.1
10
10
10
10
1021
20
FLOPs
FLOPs
FLOPs
Figure Investigation of the optimal number experts for three different model sizes, 2B, 5B, and 10B; and in three different scenarios, from left right: simply measuring the size of the model, including the size of KV-cache with 32k tokens, and including the inference cost of processing 100B tokens.
at smaller scales than in simple compute-optimality due to decreased inference FLOPs. The E = and E models shown in the Figure use respectively 36% and 61 % less FLOPs per token than their dense counterpart.
Rule of Thumb. An MoE model with E < ex- perts, trained on E-times more tokens than compute- optimal dense model, outperforms it while maintaining the same total parameter count.
4.5. Summary
The notions of inference optimality and total memory op- timality can naturally be combined. Figure 3 (c) presents comparison between different numbers of experts, where the KV cache is included in the model's memory require- ments and the compute budget is shared between training and inference. Finally, Figure investigates the optimal E for a sample of model sizes while including the KV cache and considering the inference cost.
Note that, in this scenario FLOPs matched MoE will gen- erally have less than E-times larger dataset, but we wanted to keep this rule simple and conservative. Detailed compar- isons and differences between memory and FLOPss matched models can be found on Figures & 4.
It is important to remember that such scaling might not always be possible practice due to limited dataset sizes. This points towards a possible drawback while using MoE and underscores the need for growth in dataset sizes.
For practitioners, as simplification of our analysis, we propose general rule of thumb:
6

---

Joint Scaling Laws for Mixture of Experts
3.6
38
Experts
Training Thkorrss
3.4
3.6
A5
9.2
3.4
12.00
Predicted
1drll Validation
Lo00
3.2
2.8
2.0
288
24
28
028
D:4
0.00
1:00
Observed Loss
Active Paramrs
Activo Paraarss
(a)
(b)
(C)
Figure 5. (a) Quality of the fit. The maximum absolute error on the held-out extrapolation is 0.018. (b) Predicted loss compared with an observed loss for E 1. (c) Predicted loss (dashed line) compared with observed loss for E 4. We that the training dataset, the error increases in an undertrained setting (D/N < more parameters than tokens). However, this scenario is never practical from our perspective.
Transformer model with Switch (Fedus et al., 2022) layers, using standard values of router z-loss 0.001 and load balanc- ing loss 0.01. The GPT-2 tokenizer (Radford al., 2018) is employed. For better stability, weight initialization follows truncated normal distribution with reduced scale of 0.1, as suggested by (Fedus et al., 2022). Mixed precision training is used, with the attention mechanism, position embeddings RoPE (Su et al., 2024) and router always maintained at high precision. The models use the SwiGLU activation (Shazeer, 2020) with hidden size equal 3d and activate one expert per token (unless the token is dropped due limited capacity). For evaluation, we increase the capacity factor to ensure dropless processing of the tokens.
Table Optimal E for different training budgets and three typical memory constraints, corresponding GPU, and an 8xH100 GPU node. We assume 16k tokens in the KV cache and bfloat16 for storing model weights and activations.
RTX4090 GPU, H100
Memory Constraint
Training Compute
24GB
80GB
640GB
1 x 10
16
2 32
2 32
1 x 1022
16
2 32
1 X 1023
8
2 32
1 X 1024
5. Fitting the Scaling Law
5. BATCH SIZE RAMP-UP
In this section, we present details of experiments and the procedure of fitting the scaling law parameters, see Table 3 in the Appendix. Those results are based on extensive, large- scale empirical evidence, including 270 models with up to 5B parameters, trained on variety of compute budgets. For a full list of experiments, see Appendix H.
Performance of a deep learning optimization procedure can suffer as result of using an exceedingly large batch size (McCandlish et al., 2018). To mitigate this potential issue, especially early in the training, ramp-up. Similar strategies training runs (Rae et al., 2022; Dubey et al., 2024). We in- crease the batch size from 64K to 128K after 0.5B training tokens and further to 256K after 1B training tokens. Instead of using noise scale as a critical batch size predictor (Mc- Candlish al., 2018) we opted for straightforward grid directly predict a transition point after which increased batch size does not impair performance.
employ batch-size used in contemporary LLM
5.1. Model Hyperparameters
The selection of hyperparameters and training details is crucial for ensuring the robustness of scaling laws (Porian et al., 2024; Pearce Song, 2024). In our work, we employ of best practices and modern design choices, aiming to provide accurate predictions applicable to real-life practice.
1.1.2. LEARNING RATE SCALING
All models used in this study are decoder-only Transformers trained on the highly filtered FineWeb-Edu (Penedo et al., 2024). is a subset of Web, whose curation process was guided using popular benchmarks. FineWeb-Edl is se- lected using a filter for highly educational content. We use
Kaplan et al. (2020) have shown that scaling laws for hyper- parameters can be used to adjust them according to the size of the model in the case dense Transformers. For MoE models, find the literature inconclusive-while some (Dai

---

Joint Scaling Laws for Mixture of Experts
al., 2024) pretrain MoËs with lower LR than correspond- ing dense models, others (Zoph al., 2022) report better performance when finetuning MoEs with higher learning rates. To fill this gap. we derive scaling law for the peak learning rate for MoE based on the number of active non- embedding parameters Nactre and the number of experts E:
Further details about our scaling rule for learning rates can be found in the plots in Appendix G.
5.1.3. LEARNING RATE SCHEDULE
Hägele et al. (2024) suggest that a constant learning rate schedule can yield similar performance to other established methods, such as the cosine schedule. At the same time, it offers a valuable advantage when varying training duration, as intermediate checkpoints can be reused when training models for longer time. With cosine schedule, interme- diate checkpoints can introduce bias into the fit, according to the analysis of Kaplan et al. (2020) by Hoffmann et al. (2022). We employ a constant learning rate schedule with a linear warmup over the initial 130M tokens and with a linear decay from the peak learning rate to 0 over the final 20% of tokens. For each model size, longer runs reuse intermediate checkpoints from the shorter ones.
LR(Nactre: B) exp
-0.25 In( E));
(10)
and use this equation to set the learning rate in our main scaling laws experiments. We fit the coefficients of this equation using the least squares method, minimizing the error between the prediction and the optimal learning rate from the experiment grid. Contrary to Kaplan et al. (2020), use linear transformation of the parameter count to predict the logarithm of the learning rate, instead of directly predicting the learning rate. This approach allows us to avoid the breakdown of the formula above 10 parameters mentioned in their work, where the predicted learning rate becomes negative. This phenomenon is independent of the actual fit and is simply a property of the formula used. Besides being well-defined in the extrapolation, we argue that optimal learning rates visibly follow this logarithmic trend, as seen in Figure 8 in Appendix.
5.2. Optimization of Formula Coefficients
Following Hoffmann et al. (2022), we use the LBFGS al- gorithm to optimize the coefficients of formula 6. See Appendix B for details. We observe good fit with RMSE, = 0.0039 on a held-out set of our 30 runs with the lowest loss, and RMSE, = 0.0062 on the training dataset. To further verify the validity of our formula, we train sep- arate Chinchilla scaling laws for different E using the same hyperparameters and the corresponding subset of the initializations grid. This approach serves as lower bound for loss of our joint formula on the training dataset, as it can emulate its coefficients; however, it is more prone to overfitting because effectively more parameters are utilized, Using this approach, we obtain lower error on the training dataset of RMSE?" = 0.0059 and marginally higher on the 0.0041. We believe this is strong confirmation that our joint formula is actually describing how variable E influences training.
Finding 4. More experts + lower learning rate, Increasing the number of experts in MoE model should be accompanied by lowering the learning rate accord- ingly (Figure in Appendix).
The second difference between our formula and the one by Kaplan al. (2020) is incorporating the number of experts, allowing us to model the optimal behavior of this hyperpa- rameter across dense models and different MoEs. This is an important detail that allows unbiased comparison among different models, ensuring that each one is optimally tuned. Furthermore, it allows us to answer the question of whether MoE should be trained with a lower or higher LR. While our formula accommodates both scenarios, we can clearly see in Figure 8 in Appendix that increasing E requires lower learn- ing rates, resulting in negative value for the coefficient. Moreover, we verify this thesis by tuning the fit on E = and F. = 8, and validating it on interpolation E = and extrapolation E = 32. In both cases, the validation predicts the optimal learning rate for the model configuration or a value with practically the same performance.
validation RMSE
In Figure 5, we visually verify the extrapolation of the joint fit. Prediction errors are categorized by different numbers of experts, highlighting that our joint formula is not biased for any specific E.
6. Limitations and Future Work
In our work, we focus on the standard MoE variant, where the size of the expert is the same as the size of the feed- forward layer of a corresponding dense model. Some re- cent findings (Dai et al., 2024; Ludziejewski et al., 2024; Muennighoff et al., 2024; Team, 2024b) indicate that fine- grained MoE models would enhance our reported benefits of using MoE. Simi- larly, adopting a dropless MoE (Gale et al., 2023) approach instead of relying on a capacity factor could lead to fur-
In Figure 9 in Appendix, we perform an ablation of this additional power law on E by repeating our entire fitting procedure without the E' component. This shows, especially with the extrapolation on E = 32, that dependence on E is crucial, and its omission can impair the performance of MoEs.
more efficient and, most probably,
8

---

Joint Scaling Laws for Mixture of Experts
ther improvements. We leave the integration of those MoE improvements for future work.
computational resources provided by Writer.
Moreover, our Chinchilla-basse optimality analysis uses FLOPs, that may not reflect wall-clock training time of models with different architectures. While analyzing to- tal parameter, instead active parameter matched models partly alleviates this issue because of the same memory- bottleneck, various implementations and distributed training algorithms are not considered in this work.
Impact Statement
Recent rapid advancements in language models have in- troduced numerous new capabilities while simultaneously sparking concerns about their societal impact and prompt- ing intense discussions within the community. Our work enhances the understanding of scaling laws and offers guid- ance on improving the efficiency of LLMs, potentially am- plifying existing risks by extending LLM capabilities. How- ever, our approach does not introduce any novel threats. Therefore, we refer to the existing literature on the broader impact of language models (e.g., Borgeaud et al. (2022)).
We assumed, the Chinchilla scaling law (2) as the basis of our formulas. While this is well-grounded in literature, this formula is known to have limitations, especially for a wide range of ooken-to-parament ratios. We observed this also in some of our experiments, as outliers often are highly under or over-trained.
References
Borgeaud, S., Mensch, A., Hoffmann, J., Cai, T., Rutherford, E., Millican, K., Van Den Driessche, G. B., Lespiau, J.-B., Damoc, B., Clark, A., et al. Improving language models by retrieving from trillions of tokens. In International conference on machine learning, pp. 2206-2240 PMLR, 2022.
7. Conclusions
In this work, we derived the joint scaling laws for Mixture of Experts, relating the loss of the model to the number of parameters, the number of training tokens, and the num- ber of experts. By considering both compute and memory constraints, as well as the expected inference workload, we demonstrated that MoE models can outperform dense models even when constrained by memory usage or total parameters, contrary to common assumptions and intuitions that MoE models are more memory-intension than dense models.
Clark, A., de las Casas, D., Guy, A., Mensch, A., Paganini, M., Hoffmann, J., Damoc, B., Hechtman, B., Cai, T. Borgeaud, S., van den Driessche, G., Rutherford, E., Hen- nigan, T., Johnson, M., Millican, K., Cassirer, A., Jones, C., Buchatskaya, E., Budden, D., Sifre. L., Osindero, S., Vinyals, O., Rae, J., Elsen, E., Kavukcuoglu, K., and Simonyan, K. Unified scaling laws for routed language models, 2022.
Our analysis reveals how the optimal training strategies shift as the number of experts varies. This provides principled framework for selecting MoE hyperparameters under given constraints, highlighting the trade-offs between memory and compute performance.
Dai, D., Deng, C., Zhao, C., Xu, R. X., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., Xie, Z., Li, Y. K., Huang, P., Luo, F,, Ruan, C., Sui, Z., and Liang, W. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models, 2024.
Acknowledgments
We would like to express sincere gratitude to Szymon Anto- niak and Piotr Padlewski for their detailed comments and invaluable discussions. We also thank Konrad Staniszewski for his feedback on the draft of this paper.
DeepSeek-AL, Guo, D., Yang, D.,. Zhang, H., Song, J., Zhang. R., Xu, R., Zhu, Q., Ma, S., Wang. P., Bi, X., Zhang, X., Yu, X., Wu, Y., Wu, Z. F., Gou, Z.,, Shao, Z., Li, Z., Gao, Z., Liu, A., Xue, B., Wang, B., Wu, B., Feng, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., Dai, D., Chen, D., Ji, D., Li, E., Lin, F., Dai, F., Luo, F., Hao, G., Chen, G., Li, G., Zhang, H., Bao, H., Xu, H., Wang, H., Ding, H., Xin, H., Gao, H., Qu, H., Li, H., Guo. J., Li, J., Wang, J., Chen, J., Yuan, J., Qiu. J.. Li, J., Cai, L., Ni, J., Liang, J., Chen, J., Dong, K., Hu, K., Gao, K., Guan, K., Huang, K., Yu, K., Wang. L., Zhang, I.,, Zhao, I.., Wang, I.,, Zhang. L., Xu, L., Xia, L., Zhang, M., Zhang, M., Tang, M., Li. M., Wang, M., Li, M., Tian, N., Huang, P., Zhang, P., Wang, Q., Chen, Q., Du, Q. Ge. R., Zhang, R., Pan, R., Wang, R., Chen, J., Jin, R. L.., Chen, R., Lu, S., Zhou, S., Chen, S., Ye,
We gratefully acknowledge the Polish high bigh-performance computing infrastructure PLGrid (HPC Center: ACK Cyfronet AGH) for providing computer facilities and sup- port within the computational grant no. PL.G/2224/017066 This research was partially supported by the Na- tional Science Centre (Poland) grant Sonata Bis 2024454/E/SST6/00388, the National Science Centre (NCN) Grant no. 2020/37/B/SST6044199 the National Science Centre (NCN) Preludium Grant no. 2022/45/N/ST6/02222. the "European Lighthouse of AI for Sustainability" ELIAS grant no. 101 100122237, and the NCBiR grant Part of the experiments utilized
9

---

Joint Scaling Laws for Mixture of Experts
S., Wang, S., Yu, S., Zhou, S., Pan, S., Li, S. S., Zhou, S., Wu, S., Ye, S., Yun, T., Pei, T., Sun, T., Wang, T., Zeng, W., Zhao, W., Liu, W., Liang, W., Gao, W., Yu, W., Zhang, W. Xiao, W. L., An, W., Liu, X., Wang, X.. Chen,
L., Tang, E., Thite, A., Wang, B., Wang, K., and Zou,
A. A framework for few-shot language model evaluation,
12 2023, URL https: http///eoodd.//// orggrecorda 10256836.
X. Nie, X., Cheng, X., Liu, X., Xie, X., Liu, X., Yang, X., Li, X., Su, X., Lin, X., Li, X. Q. Jin, X., Shen, X., Chen, X., Sun, X., Wang, X., Song, X., Zhou, X., Wang,
Ghorbani, B., Firat, O., Freitag, M., Bapna, A., Krikun, M., Garcia, X., Chelba, C., and Cherry, C. Scaling laws for neural machine translation, 2021.
X. Shan, X., Li, Y. K., Wang. Y. Q. Wei, Y. X., Zhang, Y., Xu, Y., Li, Y., Zhao, Y., Sun, Y. Wang, Y., Yu, Y., Zhang, Y. Shi, Y., Xiong, Y., He, Y., Piao, Y. Wang, Y., Tan, Y., Ma, Y., Liu, Y., Guo, Y., Ou, Y. Wang, Y., Gong, Zou, Y., He, Y. Xiong, Y., Luo. Y. You, Y., Liu, Y., Zhou, Y., Zhu, Y. X., Xu, Y., Huang, Y., Li, Y.. Zheng, Y., Zhu, Y., Ma, Y., Tang, Y., Zha, Y., Yan, Y., Ren, Z. Z.,, Ren, Z., Sha, Z., Fu, Z., Xu, Z., Xie, Z., Zhang, Z., Hao, Z. Ma. Z. Yan, Z., Wu, Z., Gu, Z., Zhu, Z.,, Liu, Z., Li, Z., Xie, Z., Song, Z., Pan, Z., Huang, Z., Xu, Z.. Zhang, 7., and Zhang, 7. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning, 2025. URL https: arxiv. org/ abs/ 2501 12948.
Hägele, A., Bakouch, E., Kosson, A., Allal, L. B., Werra, L. V., and Jaggi, M. Scaling Laws and Compute-Opting Training Beyond Fixed Training Durations. Advances in Neural Information Processing Systems, 2024. URL http://arxii org/ www./b/4000 18392.
He, H. Making deep learning go ьиит from first principles. 2022. URL hettps://horae.... html.
intro.
Henighan, T., Kaplan, J., Katz, M., Chen, M., Hesse, C., Jackson, J., Jun, H., Brown, T. B., Dhariwal, P., Gray, S., Hallacy, C., Mann, B., Radford, A., Ramesh, A., Ryder, N., Ziegler, D. M., Schulman, J., Amodei, D., and Mc- Candlish, S. Scaling laws for autoregressive generative modeling. 2020.
Du, N., Huang, Y., Dai, A. M., Tong, S., Lepikhin, D., Xu, Y., Krikun, M., Zhou, Y., Yu, A. W., Firat, O., Zoph, B., Fedus, L., Bosma, M., Zhou, Z., Wang. T., Wang, Y. E., Webster, K., Pellat, M., Robinson, K., Meier-Hellstem, K., Duke, T., Dixon, L., Zhang, K., Le. Q. V., Wu, Y., Chen, Z., and Cui, C. Glam: Efficient scaling of language models with mixture-oofeexperts 2022.
Hestness, J., Narang, S., Ardalani, N., Diamos, G., Jun, H., Kianinejad, H., Patwary, M. M. A., Yang, Y., and Zhou, Y. Deep learning scaling is predictable, empirically, 2017.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21788, 2024.
Hochreiter, S. and Schmidhuber, J. Long short-term memory.
Neural computation, 9(8): 1735-17758788 1997.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Rae, J. W., Vinyals, O., and Sifre, L. Training compute-optimal large language models, 2022.
Fedus, W., Zoph, B., and Shazeer, N. Switch transform- ers: Scaling to trillion parameter models with simple and efficient sparsity, 2022.
Frantar, E., Riquelme, C., Houlsby, N., Alistarh, D., and Evei, U. Scaling laws for sparsely-connected foundation models, 2023.
Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. E. Adaptive mixtures of local experts. Neural Computation, 3(1):79-87, 1991. doi: 10. 10.1162/ncco. (tt////cco.199...111.. 1991.3.1.79.
Gale, T. Narayanan, D., Young. C., and Zaharia, M. Megablocks: Efficient sparse training with mixture-of-experts. In Song, D., Carbin, M., and Chen, T. (eds.), Proceedings of Machine Learning and Systems, volume 5, pp. 288-304. Curan, 2023. URL https://proceecingg http//apee_fillls/pppp////33301 2aS=f1993766efffcee2222ccccce pdf.
Jiang, A. Q. Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., de las Casas, D., Hanna, E. B., Bressand, F., Lengyel, G., Bour, G., Lample, G., Lavaud, L. R., Saulnier, L., Lachaux, M.-A., Stock, P., Subramanian, S., Yang, S., Antoniak, S., Scao, T. L., wet--mmmseeeppe Geone, 3r,, Lavril, T., Wang. T., Lacroix, T,, and Sayed, W. E. Mixtral of experts. 2024.
Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi,
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models, 2020.
A., Foster, C., Golding, L., Hsu, J., Le Noac'h, h, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang,
J., Reynolds, L.., Schoelkopf, H., Skowron, A., Sutawika,
10

---

Song, F., Aslanides, J., Henderson, S., Ring. R., Young, S. Rutherford, E., Hennigan, T., Menick, J., Cassirer, Joint Scaling Laws for Mixture of Experts
Kumar, T., Ankner, Z., Spector, B. F., Bordelon, B., Muen- nighoff, N., Paul, M., Pehlevan, C., Re, C., and Raghu- nathan, A. Scaling laws for precision. In The Thirteenth International Conference on Learning Representations, 2024.
Rae, J. W., Borgeaud, S., Cai, T., Millican, K., Hoffmann, J.,
A., Powell, R., van den Driessche, G., Hendricks, L. A., Rauh, M., Huang, P.-S., Glaese, A., WelbI, J., Dathathri, S., Huang, S., Uesato, J., Mellor, J., Higgins, I., Creswell, A., McAleese, N., Wu, A., Elsen, E., Jayakumar, S., Buchatskaya, E.. Budden, D., Sutherland, E., Simonyan, K., Paganini, M., Sifre, L.., Martens, L.., Li, X. L.., Kun- соrо, A., Nematzadeh, A., Gribovskaya, E., Donato, D., Lazaridou, A., Mensch, A., Lespiau, J.-B., Tsimpoukelli, M., Grigorev, N., Fritz, D., Sottiaux, T., Pajarskas, M., Pohlen, T., Gong, Z., Toyama, D., de Masson d' Autume, C., Li, Y., Terzi, T., Mikulik, V., Babuschkin, I., Clark, A., de Las Casas, D., Guy, A.. Jones, C., Bradbury, J., Johnson, M., Hechtman, B., Weidinger, L., Gabriel, I., Isaac, W., Lockhart, E.,, Osindero, S., Rimell, L., Dyer. C., Vinyals, O., Ayoub, K., Stanway, J., Bennett, I.., Hassabis, D., Kavukcuoglu, K., and Irving, G. Scaling language models: Methods, analysis and insights from training gopher, 2022.
Langley. P. Crafting papers on machine learning. In Langley, P. (ed.). Proceedings of the 17th International Conference on Machine Learning (ICML 2000), pp- 1207-1216, Stanford, CA, 2000. Morgan Kaufmann.
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y.,
Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models with conditional computation and automatic sharding, 2020.
Ludziejewski, J., Krajewski, J., Adamczewski, K., Pióro, M., Krutul, M., Antoniak, S., Ciebiera, K., Król, K., Odrzygóádú, T., Sankowski, P., Cygan, M., and Jaszczur, S. Scaling laws for fine-grained mixture of experts. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.). Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 33270-33288, PMLR, 21-27 Jul 2024. URL hitps: /ttpp/edccceeee...cco mlr. press/ v235 112255/Ludzieejwss.see
Riquelme, C., Puigcerver, Mustafa, B., Neumann, M., Jenatton, R., Susano Pinto, A., Keysers, D., and Houlsby, N. Scaling vision with sparse mix- ture of experts. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances Neural Information Processing Systems, volume 34, pp.
html.
McCandlish, S., Kaplan, J., Amodei, D., and Team, O. D. An empirical model of large-batch training, 2018. URL https: //tt:////rxivv// arxiv. org/ abs///b///// 11///11111.1 06162.
Curran Associates, Inc., 2021. URL https:: neurips. cc /ttpppper_filleeppperr////// 48 1823749f22ea0c7442222233366 pdf.
Muennighoff. N., Soldaini, L., Groeneveld, D., Lo, K., Mor- rison, J., Min, S., Shi, W., Walsh, P., Tafjord, O., Lambert, N., Gu, Y., Arora, S., Bhagia, A., Schwenk, D., Wadden, D., Wettig, A., Hui, B., Dettmers, T., Kiela, D., Farhadi, A., Smith, N. A., Koh, P. W., Singh, A., and Hajishirzi. H. Olmoe: Open ofoxxuur--of-epertt language models, 2024. URL https: / /aaaiiiiiii org/ abs / 2409.4440 02060.
Paper.
Sardana, N., Portes, J., Doubov, S., and Frankle, J. Be- yond chinchilla-optimal:: accounting for inference in language model scaling laws. In Proceedings of the 41st International Conference on Machine Learning, pp. 43445 43460, 2024.
Pearce, T. and Song, J. Reconciling kaplan and chin- chilla scaling laws. Transactions on Machine Learning Research, 2024.
Shazeer, N. Glu variants improve transformer, 2020. URL https: https://arxiv..orr/aae///000 05202.
Penedo, G., Kydličck, H., Lozhkov, A., Mitchell, M., Raf- fel, C. A., Von Werra, L., Wolf, T., et al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:30811-308A4 1--0099, 2024.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. Outragcously large neural networks: The sparsely-gate mixture-of-experts layer, 2017.
Porian, T., Wortsman, M., Jitsev, J., Schmidt, L., and Car- mon, Y. Resolving discrepancies in compute-optimal scal- ing of language models. Advances in Neural Information Processing Systems, 37:100535-000570, 2024.
Shazeer, N., Cheng, Y., Parmar, N., Tran, D., Vaswani, A., Koanantakool, P., Hawkins, P., Lee, H., Hong, M., Young. C., Sepassi, R., and Hechtman, B. Mesh-tensorflow: Deep learning for supercomputers, 2018.
Radford, A., Narasimhan, K., Salimans, T., and Sutskever, Improving language understanding by generative pre- training. 2018.
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 2024.

---

Joint Scaling Laws for Mixture of Experts
Team, Q. Qwen2.5 technical report. arXiv preprint 1rXiv:2412.55115, 2024a.
Team, Q. @wen1.5-moc. Matching 7b model per-
formance with 1/3 activated parameters", February 2024b. URL https://qwenli: github. -io/ blog/qwent moog///
Yun, L., Zhuang. Y., Fu, Y. Xing, E. P., and Zhang, H. Toward inference-optimal mixture-of-exper large lan- guage models, 2024. URL attps:://rxiis abs/2404 02852.
Zadouri, T., Üstün, A., Ahmadian, A., Ermis, B., Locatelli, A., and Hooker, Pushing mixture of experts to the limit: Extremely parameter efficient moe for instruc- tion tuning. In The Twelfth International Conference Learning Representations, 2024.
Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. Scaling vision transformers.
Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12104 12113, 2022.
Zoph, B., Bello, I., Kumar, S., Du, N., Huang, Y., Dean, J., Shazeer, N., and Fedus, W. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08906, 2022.
12

---

Joint Scaling Laws for Mixture of Experts
A. Technical Details
A.1. Counting Parameters
There are many ways the size model can measured. The two most important distinctions are whether total or active parameters are counted and whether the parameters in the embedding and unembedding layers are counted. Various papers assume different notations, notably Kaplan al. (2020) use nonembedding parameters while Hoffmann et al. (2022) opt for the parameter count including embedding and unembedding. Throughout work, try to make it clear which way of counting we are using in each particular instance. When additional information is given, Natt and Notal denote respectively active and total parameters, including the embedding and unembedding.
we let dmnndel be the hidden dimension of model, and dvacab be the vocabulary size (50,257 in our case), then the following relations hold:
Numil wraahh + (4 + N & Nat = md + 13.Nttocks montel
(11)
(12)
A.2. Counting FLOPs
Basing on Sardana et al. (2024), we assume the cost of training to be Piaiiinng = as Dimining and the cost of inference to be Finference 2/ 2n Due to the relatively small number (S. 32) of experts used with implicit expert granularity of (Ludziejewski al., 2024), can consider the memory and FLOPs of routing to be negligible, following Clark et al. (2022).
A.3. Model Configs
The vast majority of our experiments simple rule for scaling the config, i.e. Nbtocks Nucads dmoodl/ /4 and assume these relations hold in all calculations. We base this rule findings by Kaplan et al. (2020).
B. Fit Details
Table 3. Fitted coefficients of our joined formula.
a
8
w
C
Extart
Emax
C
35.91
-0.1889
-0.2285
0.0098
35.98
-0.1775
0.5529
-0.0259
2.0732
290.4521
1.3637
Table 4. The fitted coefficients of our joint formula, Equation (6), reduced to the Chinchilla scaling law. Equation (2), for given number experts, E. We observe that the dataset exponent, v. increases significantly. This one the reasons why compute-optimal parameter-to-token ratios change with E.
m
/
n
V
30.3640
-0.1817
53.9838
-0.1965
1.3637
2
27.7982
-0.1780
66.8401
-0.2065
1.3637
4
24.8462
-0.1731
87.7022
-0.2192
1.3637
8
21.8330
-0.1676
119.9126
-0.2338
1.3637
16
19.0159
-0.1617
167.5073
-0.2494
1.3637
32
16.5424
-0.1557
234.6726
-0.2652
1.3637
Following Hoffmann et al. (2022), use the LBFGS algorithm with learning rate of 1c-4 and weight decay of 1e-5 to fit the coefficients of Equation 6, optimizing the Huber loss with = 0.01 over the set of our training runs described in table in Appendix H. Instead of removing outliers and underperforming models from the training set, we underweight them proportionally to the loss. Optimization hyperparameters were manually tuned to minimize error over the training dataset. The final fitted coefficients of Equation are within the boundaries of the grid of initializations given by: a (0.05,0) 0.25, 0.5), B E (0.05, 0.25, 0.5). A E (30, 100, 300), B E {30, 100, 300). C E [0.5, 1,2),
13

---

Joint Scaling Laws for Mixture of Experts
d E (,,55, E (,05,, w E (-0.5, 0.53, < (-0.5, 0.5). The selected coefficients were those with the lowest score, defined as the of RMSE the training and held-out extrapolation validation set. The formula in Equation 6 was calculated in logarithm, without any exponentials, using only linear transformations and the logsumexp operation. It was optimized to predict the logarithm of L, and parameters b. and e were optimized in logarithm. All these steps were taken to increase numerical stability and were essential for proper convergence.
C. Derivation of Nopt and Dopt act
To derive the optimal aets D given some compute budget F and E, one needs solve:
arg min Lg (Nacts D) s.t. F Nact D. Naa.
To solve for Nact, substitute: and set the derivative to O:
D F
6.N
d
Lg (Nact, D)
d d.Vact
m (EE Na((e) act
+ n(B) (E)
v(E)
v-v(E)
0.
d.
act
After rearranging:
Noptt (m(E) (m())u(( ((() (m MMM)
vi
((
(((E)v())
The derivation of Dept is analogous.
D. Token Dropping Analysis
In general, observe that the load balancing loss quickly induces balance between experts. Overall, the percentage of dropped tokens is low and doesn't it exceed 10%, therefore doesn't significantly affect the training efficiency. Below present plot of pay-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 observe relatively the largest ratios dropped tokens in initial and last layers.
Fraction of Dropped Tokens
Dropped dmodel= 1024,
0.08
Number Experts - 2.0
0.04
10.0
Block
Fraction of Dropped Tokens
Dropped tokens for dmodel=768, blocks
Number Experts
2.0
0.08
0.06
16.0
0.04
0.02
Block
Figure Dropped tokens for selected models.
14

---

Joint Scaling Laws for Mixture of Experts
E. Downstream Performance
In addition to measuring pretraining loss, we evaluate downstream performance using LM Evaluation Harness (Gao et al., 2023). Unlike in training, we employ dropless MoE. We observe a strong correlation between the pretraining perplexity and downstream performance across all E"s. The results can be seen in Figure 7. On some tasks (HellaSwag, Winogrande, SciQ), dense models seem to outperform MoE models given the same pretraining perplexity. They also seem to be more robust to domain shift in language modeling, exemplified by the results on the LAMBADA benchmark. On some other tasks (e.g. Open MoE models seem to fare similarly or slightly better than dense models if they have been trained to the same pretraining perplexity.
0.7
D.
Acarsoy thane
Accaaay thome
2.55
0.3
2.45
0.25
Fandem dharce 026
0.4
0.6
0.55
Accunay (War))
0.5
0.45
088
03
055
0.6
rogurds
Ope
0.58
0.56
0.4
0.35
2:48
AQA
0.55
0.22
AccU (Warm)
0.8
0.8
3.55
066
0:35
chonce
Rondort chance 0:88
0:55
LAVSADA Skankad
LAMBADA--SA
0.4
0.3
0.25
0.
3.
Faas Looe
No. Exparts
Figure 7. Downstream performance of the models.
15

---

Joint Scaling Laws for Mixture of Experts
F. Bootstrap Results
To quantify the uncertainty of our derived results, we calculate bootstrapped results for the optimal number active parameters and training tokens, following (Hoffmann al., 2022). The results are shown in Table We sample 80% of data 100 times and report the 10th and 90th percentiles.
Table 5. Bootstrap intervals for the optimal No and D across training budgets and expert counts.
Training Budget
Experts
Nept
(90% interval) act
Dopt (90% % interval)
1 x 10
1.5B B.1B
7.9B-11.1B
2
1.4B-1.8B
9.1B-12.1B
4
1.2B-1.6B
1.3B-14.1B
8
940.0M-1.5B
11.4B-17.7B
16
732.7M-1.3B
12.6B-22.8B
559.5M-1.2B
13.7B-29.88
1 x 10
1
4.5B-7.8B
21.4B-33
2
4.1B-7.3B
22.7B-40.9B
3.5B-7.0B
23.7B-47,
2.8B-6.8B
24.7B-60.1B
16
2.0B 6.6B
32
1.6B-6.5B
25.7B-104.1B
x 10
-3.BB-00.0B
55.5B-125.38
12.1B-30.0B
55.6B-137.3B
10.3B-30.6B
54.5B-161.4B
8
8.0B-31.5B
16
5.9B-32.9B
50.6B-284.4D
32
4.4B-34.5B
48.4B-380..
16

---

Joint Scaling Laws for Mixture of Experts
G. Learning Rate Scaling Fit
2ep.
2.AAR Locun
4732"
B50 WE 60-31.00
RST 29 18, anvin andiiod
76, 01:30
B1 1. 2i.
D 1. sn
pockit
14"11
12x
33107
Figure& 8. Visualization of the fit (E E (1,8)) of our LR scaling rule, interpolation (E 1) and extrapolation (E=32) 32).
Learming (alem))
rateml)
Leatring Nata rale-
Seariing ratem321
Pola 390.09
4'38
100.0
380 106
150
21.35.6
45.00
LCA 40.00
43
40,00
31x1F
16x12'
32>13"
34307
SANTing
CAST
LANTTIN 1ore
Figure 9. Ablation for the LR scaling rule fit without considering the number of experts E. While performance on the training set (E {1,8)) looks acceptable, the extrapolation E 32 is clearly suboptimal, validating the need for considering E.
17

---

Joint Scaling Laws for Mixture of Experts
H. Experiments Listing
Noval
_hads heads
blocks
desodel
Nac
E
D
5.0B
16
16
1024
321M
32
16.0B, 8.0B, 4.0B, 2.0B, 1.0B, 500M
3.8B
28
28
1792
1.3B
4
11.1B, 5.6B, 2.8B, 2.0B
3,3B
11
21
1408
683M
8
16.0B, 8.0B, 4.0B, 2.0B, 1.0,, 500M
3.0B
26
26
1664
1.1B
4
80.0B, 64.0B, 48.0B, 32.0B, 16.0B, 8.0B, 4.0B, 2.0B, 1.0B, 500M
2.7B
36
36
2304
2.7B
1
9.2B, 5.5B, 2.8B, 2.0B, 1.4B, 980M
2.6B
30
30
1920
1.6B
2
5.4B, 2.7B
2.6B
1024
16.0B, 8.0B, 4.0B, 2.0B, 1.0B,
2.2B
28
28
1792
1.3B
2
18.6B, 11.I., 5.6B. 4,0B, 2.8B, 2.0B
2.1B
12
12
768
169M
32
8.0B, 4.0B, 2.0B, 1.0B, 500M
2.1B
10
1280
469M
8
32.0B, 16.0B. 8.0B, 4.0B, 2.0B, 1.0B
1,9B
22
22
1408
709M
35.3B, 12.2B, 10.6B, 7.7B, 5.3B, 3.8B
1.8B
11
21
683M
8.0B, 16.0B, 4.0B, 2.0B, 1.0B, 500M
1.8B
26
26
1664
1.1B
2
16.0B, 8.0B, 4.08, 2.08, 1.0B, 500M
1.6B
30
30
1920
1.6B
5.4B, 2.7B
1.4B
16
1024
321M
8
16.0B. 8.0B, 4.0B, 2.0B, L.B,, 500M
1.3
28
1.3B
6.5B, 3.3B, 18.6B, 11.IB, 5.6B, 4.0B, 2.8B, 2.0B
1.3B
640
118M
32
4.0B, 2.0B, 1.0B, 500M
1.2B
10
16
1280
469M
4
32.0B, 16.0B, 8.0B, 4.0B, 2.0B, 500M
1.1B
12
12
768
169M
16
8.0B, 4.0B, 2.0B, 1.0B, 500M
1.I
26
26
1664
1.IB
14.0B, 12.0B, 10.0B, 80.0B, 64.0B, 48.0B, 32.0B
1.B
1664
1.1B
16.0B, 8.0B, 4.0B, 2.0B, 1.0B, 500M
1.1B
22
22
1408
709M
2
3.8B, 49.8B, 24.9B, 12.5B, 6.2B, 3.IB, 1.6B, 778M
LIB
22
22
1408
709M
2
21.8B, 18.7B. 15.6B. 35.3., 12.2B, 10.6B, 7.7B, 5.3B
LIB
18
18
1152
426M
4
31.B, 25.9B. 20.7B. 10.4B, 5.2B, 2.6B. 1.3B
LIB
683M
32.0B, 16.0B.
4.0B, 2.0B, 1.0B, 500M
890M
1536
890M
9.9B,
850M
20
20
1280
555M
2
16.0B, 8.0B
774M
16
16
1024
321M
4
16.08, 8.0B, 4.0B, 2.0B, 1.0B, 500M
709M
22
22
1408
709M
35.3B, 12.2B. 10.6B. 7.7B. 5.3B. 3.8B. 12.5B, 6.2B
705M
1280
32.0B, 16.0B, 8.0B, 4.0B, 2.0B, 1.0B, 500M
683M
1I
21
1408
683M
32.0B, 16.0B, 8.0B, 4.0B, 2.0B, 1.0B, 500M
671M
10
10
640
118M
16
4.0B, 2.0B, 1.0B, 500M
664M
8
8
512
79M
32
2.0B, 1.0B, 500M
615M
768
169M
8.0B, 4.0B, 2.0B, 1.0B, 500M
555M
1280
16.0B, 8.0B
472M
16
16
1024
321M
2
16.0B, 8.0B, 4.0B, 2.0B, 1.0B, 500M
469M
10
16
1280
469M
32.0B, 16.0B, 8.0B, 4.0B, 2.0B, 1.0B, 500M
376M
10
640
118M
8
4.0B, 2.0B, 500M 500M
362M
512
79M
2.0B, 1.0B, 500M
360M
12
12
768
169M
4
8.0B, 4.0B, 2.0B. L.OB,
321
16
16
1024
321M
16.0B, 8.0B, 4.0B, 2.0B, 1.0B, 500M
289M
11
11
704
142M
4
4.5B, 2.3B, 1.1BB
285M
9
9
576
97M
8
3.3B, 1.7B
282M
832
201M
6.4B, 3.2B, 1.6B, 800M
233M
12
12
768
169M
2
8.0B, 4.0B, 2.0B, 1.0B, 500M
228M
10
10
640
118M
4
4.0B, 2.0B, 1.0B, 500M
211M
8
8
512
79M
8
2.0B. 1.0B, 500M
169M
768
169M
8.0B, 4.0B, 2.0B, 1.0B, 500M
154M
10
640
118M
2
4.0B, 2.0B, 1.0B, 500M
135M
8
8
512
79M
4
2.0B, 1.0B, 500M
118M
10
10
640
118M
1
4.0B, 2.0B, 1.0B, 500M
98M
8
8
512
79M
2
2.0B, 1.0B, 500M
79M
8
8
512
79M
2.0B, 1.0B, 500M
18


Reviews:
[{"review_id": "9HGXtoOnHr", "reviewer": "Reviewer_ftJ1", "summary": "The paper proposes a scaling law for mixture-of-experts and dense models, similar to the one used by Chinchilla (Hoffmann et al., 2022), but that incorporates the number of experts in the equations.\n\nThe proposed equation describing the scaling laws, is essentially the combination of Chinchilla with Clark et al. (2022), which explores the scaling laws of MoE models, but that didn't take into account the total number of training tokens $D$, only the number of active parameters $N$. The paper also presents different variations of the scaling laws that take into account not only total training cost, but also (expected) total inference cost, and memory constraints.\n\nThe authors fit the model running several expermients, with a wide range of training FLOPs (different number of tokens, different number of experts, and different model backbone sizes), and analyze the results. There are different key observations from the experimental results which are highlighted through the paper. Perhaps, the most interesting one is that MoE models can also be parameter efficient compared to dense models, which is commonly assumed to be false in the community.\n\n### Update after the rebuttal\nI thank the authors for clarifying my questions and addressing the typos that I highlighted during my review. I will keep my score and recommend that the acceptance of the paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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\".\nWe are also thankful for the detailed comments and suggestions, which we will apply in the revised manuscript version.\n\n**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.\n\n**Regarding \"Other Comments\"**\nIf we are able to fit the labels within Table 2, we will do so in the camera-ready version.\n\nRegarding 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.\n\nThank 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.", "claims_and_evidence": "The authors propose a unifying scaling law for dense and MoE models, that accurately predicts the training loss of a model given the total training budget, the total number of activated parameters, and the number of experts. \n\nThe experiments indeed show that the optimal number of experts for a given training budget depends on the memory constraints. When the training budget is high but the memory is highly constrained, it is often better to use a dense model than a MoE one. But if the training budget is low, a MoE model may be optimal, even with the same memory constraints than a dense model (see Figure 3, 4 and Table 2).", "methods_and_evaluation_criteria": "The paper follows standard practices when studying the optimal scaling of Transformer models for language modeling. The evaluation criteria (fitting the scaling formula and comparing the interpolation and extrapolation behaviour) is also standard.", "theoretical_claims": "There are no theoretical claims.", "experimental_designs_or_analyses": "The experiments are sound: the authors train a wide range of models on a large collection of data (FineWeb-edu), perform a reasonable tuning of hyperparameters (learning rate and batch size). They use standard practices to fitting the scaling laws (LBFGS, as Hoffmann et al., 2022).", "relation_to_broader_scientific_literature": "The key contributions of the paper are of high importance to the community of language modeling, deep learning architectures, and mixture-of-experts. The paper includes a quite comprehensive literature review, citing relevant papers in the context of MoEs and scaling laws for language modeling.", "essential_references_not_discussed": "The literature review focuses only on MoEs and scaling laws _for language modeling_, while the same techniques can be applied to other areas, such as computer vision. For instance, \"Scaling Vision with Sparse Mixture of Experts\" by Riquelme et al. (2021) presents MoEs for Vision Transformers (ViTs), and \"Scaling Vision Transformers\" by Zhai et al. (2021) presents scaling laws for ViTs.", "other_strengths_and_weaknesses": "I really appreciate that the authors highlight all the take-aways from the different sections, and give a general rule of thumb at the end of the paper. The paper is of great quality in my opinion. Congratulations to the authors!", "other_comments_or_suggestions": "- In Table 2, I would suggest labeling the rows as \"Training FLOPs\", and the columns as \"Maximum Memory\". This can be inferred from reading the table's caption, but it will be clearer to see if it was directly on the table, and it seems to be enough space to add these in it.\n\n- I think there's a typo in the caption of Figure 5. L346 reads \"($D/N < 1$ ---- more tokens than parameters)\". It is the other way around: $N$ is the number of parameters, and $D$ of tokens!\n\n- Could you add the legend in Figure 5b). I'm not sure what the brown curve represents (is it 16B training tokens?).", "questions_for_authors": "None.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5ge6jvynMt", "reviewer": "Reviewer_wsHW", "summary": "This work balances computational and memory constraints by deriving joint scaling laws for both Mixture-of-Experts (MoE) and dense models. The analysis shows that the optimal number of experts is closely tied to the available memory and compute budgets. Furthermore, experimental results suggest that MoE models can often outperform dense models. By transferring scaling laws from dense to MoE models, this study provides valuable insights for designing and deploying MoE architectures in large-scale training scenarios.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**Regarding the innovations, differences and challenges compared to previous work** \nWe believe that our paper delivers significant innovations not present in existing works and provides important value to the scientific community:\n- 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. \n- 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.\n- 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.\n- 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.\n\n**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. \nNotwithstanding, 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).\nWe will add the analysis of downstream performance in the camera-ready version of the paper.\n\n**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.\n\n**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.).\nFineWeb-Edu is selected using a filter for highly educational content. It “outperforms all openly accessible web-datasets on a number of educational benchmarks” [3].\nWe will clarify our choice of the dataset in the camera-ready version.\n\n**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.\n\nReferences:\n\n[1] Du et al., Understanding Emergent Abilities of Language Models from the Loss Perspective\n\n[2] Jelassi et al., Mixture of Parrots: Experts improve memorization more than reasoning\n\n[3] Penedo et al., The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale\n\n[4] Ludziejewski et al., Scaling Laws for Fine-Grained Mixture of Experts\n\n[5] Clark et al., Unified Scaling Laws for routed language models\n\n[6] Hoffmann et al., Training Compute-Optimal Large Language Models", "claims_and_evidence": "The claims are well-supported.", "methods_and_evaluation_criteria": "The proposed methods and valuations criteria almost make sense. However, here are some questions:\n1. I did not find a dedicated section in the manuscript that explicitly details the dataset selection process. Specifically, I wonder if the authors utilized established benchmarks such as the HumanEval dataset [1], which assesses the capability to generate Python functions of varying complexity. Could you please provide specific details about your dataset choices and experimental setups?\n\n2. I would like to discuss a point regarding performance metrics. While traditional scaling laws typically use final loss as the primary performance indicator, I am interested in your perspective on how loss values correlate with other metrics such as accuracy. Specifically, I would appreciate your thoughts on why lower loss values provide stronger evidence for conclusions such as \"MoE can often be the preferred alternative to dense models\" compared to alternative evaluation metrics.\n\n*[1]. Evaluating large language models trained on code. 2021.*", "theoretical_claims": "The work does not include any theoretical proofs. I believe the article would be significantly strengthened if the authors provided derivations of the optimal $N$ and $D$ (as shown in **equation (7)**) by formulating and solving the joint MoE scaling law objective function. This theoretical foundation would complement the empirical results and make the content more complete.", "experimental_designs_or_analyses": "The experimental designs are effectively validate the performance of the proposed principled framework for selecting the optimal MoE configuration and some interesting findings for comparsion between dense and MoE models.", "relation_to_broader_scientific_literature": "Recent works have shown that:\n\n1. For a fixed dataset size, as model size increases, the benefit of using an MoE diminishes. \n\n2. For a fixed model size, as the number of training tokens increases, the benefit of an MoE grows.\n\nThis paper further discuss trade-offs between the computational and memory costs of MoE models through the novel methods using joint MoE scaling laws, offering valuable insights for their design and deployment in large-scale training settings.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths:\n\n1. The paper propose a novel joint scaling law for MoE models to find optimal experts configuration, providing an interesting perspective for us.\n\n2. The topic for discussing the precise trade-offs between compute and memory efficiency is important in real-world and this work offer actionable insights for deploying MoE models in practical.\n\nWeaknesses:\n\nThe innovation of the proposed Joint MoE Scaling Law appears to be incremental. It builds upon previous work that established power-law relationships between final loss, model size, and dataset size [1], as well as recent MoE-related studies that treated the number of experts $E$ and model size $N$ as variables in formulas [2]. The main contribution seems to be incorporating dataset size $D$ as an additional variable to establish the equation between $L$, $E$, $N$, and $D$. \n\nCould you clarify the core innovations of your proposed method, as well as the differences and challenges in advancing beyond the aforementioned research? (If the authors address this concern, I would consider increasing my score.)\n\n*[1]. Training compute-optimal large language models, 2022.*\n\n*[2]. Unified scaling laws for routed language models, 2022.*", "other_comments_or_suggestions": "In the Related Work section, equation (2) appears to lack the constraint $C \\approx 6ND$. where $C$ denotes the floating-point operation count (FLOPs).", "questions_for_authors": "The paper employs a standard Switch MoE layer that routes each token to a single expert, which seems inconsistent with mainstream MoE architectures like Deepseek that activate multiple experts per token. Can the methods presented in the paper be adapted to scenarios where $K>1$ experts are activated per token?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VQ5KNeY09o", "reviewer": "Reviewer_njYw", "summary": "The authors, motivated by the Chinchilla scaling laws for large transformers and the popularity of Mixture of Experts (MoE) architectures, investigate a joint scaling law that can be applied to MoE models and dense models (when number of experts = 1). The loss of the model is related to the number of parameters, number of training tokens, and number of experts. The optimization involved minimizing this loss given a fixed compute budget F (assuming limiting factor was accelerator memory)\n\nThis work considers a standard MoE variant: a switch MoE. Tokens are routed to experts and a load-balancing loss is used. The Chinchilla-based optimality analyses assumes that computational efficiency can be measured by FLOPS: approximated by 2*N*D + 4*N*D for the forwards and backwards pass respectively, where N is the number of activated parameters and D is the number of training tokens\n\nRunning over 280 experiments, the authors fit the scaling law using least squares and share their four findings and rule of thumb.\n\nEdit: Given the authors' responses I have increased my score to an accept.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**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.\n\n**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.\n\n**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).\nOur 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.\n\n**Expert Imbalance** Thank you for raising this important question. To avoid exceeding the response length we refer to the answer given to reviewer V5XV.\n\n**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.\n\n**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].\n| E | - | 1 | 8 | 32 |\n|---|---|---|---|---|\n| estimation method/ model size | our estimate | torch | torch | torch |\n| 370M | 642.1 | 572.7 | 573.0 | 573.7 |\n| 890M | 1781.0 | 1702.2 | 1702.7 | 1704.4 |\n| 1.6B | 3261.4 | 3186.5 | 3187.3 | 3190.0 |\n\nIn a specific setup, MoE efficiency depends on the implementation, with many efficient ones [7,8] speeding up training and inference.\n\n**Distributed Training** Since our method is implementation agnostic, distributed setting should not have any impact on our conclusions.\n\n**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.\n\nReferences:\n\n[1] Kaplan et al., Scaling Laws for Neural Language Models\n\n[2] Hoffmann et al., Training Compute-Optimal Large Language Models\n\n[3] Ludziejewski et al., Scaling Laws for Fine-Grained Mixture of Experts\n\n[4] Kumar et al., Scaling Laws for Precision\n\n[5] Clark et al., Unified Scaling Laws for routed language models\n\n[6] Gadre et al. Language models scale reliably with over-training and on downstream tasks\n\n[7] Tan et al., Scattered Mixture-of-Experts Implementation\n\n[8] Zhao et al., DeepEP: an efficient expert-parallel communication library", "claims_and_evidence": "Yes, the authors do a good job at gathering convincing evidence from comprehensive experiments", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Checked the experiment settings in the appendix; generally looks good, i.e. optimised params", "relation_to_broader_scientific_literature": "This paper covers the broader scientific literature of MoE architecture (efficient/conditional compute) along with scaling laws for transformers (Hoffman et al., 2022). The key contribution is the combination of successfully extending the Chinchilla scaling law to cover MoEs.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n- The authors do a great job at visualising their results through very well-done figures in a way that effectively communicates the results\n- The summaries of the findings and rule-of-thumb is also very well communicated and makes the work easy to digest\n- Focusing on inference memory, not just training FLOPs is important\n- Range of experiments is vast and rigorous\n- Limitations section is well done \n\nWeaknesses:\n- The assumption that dataset-size is a variable I find unrealistic; empirically most will take the largest data that is available to them and treat it as a constant, varying model-size and training-steps\n- Missing how this works across different datasets, or is scaling law dataset independent and if so can can we reach this conclusion from some results\n- I feel that there is an over reliance on the Chinchilla-derived scaling law without modification (although I guess the fit has a low MAE)", "other_comments_or_suggestions": "- Sec 4.1 is missing an equation number for the optimization equation\n- A footnote could be added to explain the 6=4+2 weights for FLOPS for readers not very familiar with the literature (I had to check this myself in Hoffmann 2022)", "questions_for_authors": "- Have the authors analysed if the load-balancing loss fully mitigates the expert imbalance at scale? Does it help avoid the problem where many tokens in a batch are routed to the same expert causing poor hardware utilisation that the theoretical analysis misses? Asking if there is a chance that inference bottlenecks in the real-world render the theoretical analysis less useful?\n- Do the authors believe these results will hold for distributed training?\n- Have authors verified for a few experiments if estimated FLOPs match actual FLOPs?\n- Would optimisations like FlashAttention and MLA influence the observed scaling laws?\n- Do the authors believe any modifications should be made to the Chinchilla scaling laws that are MoE specific? I think the paper suggests no?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GiFtJ80i4m", "reviewer": "Reviewer_V5XV", "summary": "This paper introduces joint scaling laws for Mixture of Experts (MoE) and dense models, incorporating factors such as the number of active parameters, dataset size, and number of experts. The proposed scaling law captures interactions between these variables, enabling principled optimization of MoE configurations under compute and memory constraints.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n**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. \nHere (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.\n\nWe 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.\n\n**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:\n* 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.\n* We carefully adjust the learning rate depending on the number of experts and additionally tune the batch size.\n* 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.\n \nWe hope that this explanation is convincing. Should the reviewer have further concerns or suggestions regarding this topic, we would be happy to address them.\n\n**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.\n\n**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.\n\n**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.\n\n**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.\n\n**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.\n\nReferences:\n\n[1] Muennighoff et al., OLMoE: Open Mixture-of-Experts Language Models\n\n[2] Vavre et al., Llama 3 Meets MoE: Efficient Upcycling\n\n[3] Gale et al., MegaBlocks: Efficient Sparse Training with Mixture-of-Experts\n\n[4] Tan et al., Scattered Mixture-of-Experts Implementation\n\n[5] Kaplan et al., Scaling Laws for Neural Language Models\n\n[6] Clark et al., Unified Scaling Laws for Routed Language Models\n\n[7] Hoffmann et al., Training Compute-Optimal Large Language Models\n\n[8] Frantar et al. Scaling Laws for Sparsely-Connected Foundation Models \n\n[9] Ludziejewski et al., Scaling Laws for Fine-Grained MoE\n\n[10] Dai et al., DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "Related to LLM training methods.", "essential_references_not_discussed": "Yes", "other_strengths_and_weaknesses": "Strengths:\n- Provides actionable insights for MoE with different budget.\n- Extensive empirical validation (280 models, up to 5B parameters) strengthens confidence in the scaling law.\n\nWeaknesses:\n- The analysis assumes dataset size can scale freely, which may not hold in real-world scenarios with limited data.\n- Limited discussion of challenges in training MoE models (e.g., expert imbalance, routing instability) that could affect scalability.", "other_comments_or_suggestions": "It will be better to more clearly demonstrate is the comparison on MoE's full parameters or the activated parameters.", "questions_for_authors": "- How to make sure the comparison between dense and MoE is fair.\n- Would the scaling laws hold for MoE variants with different design choice, e.g., routing policy.\n- How does varying expert size (corase expert or fine expert) affect the proposed scaling law?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In the paper, the authors propose unified scaling laws for both dense and Mixture-of-Experts (MoE) models, taking into account important variables such as the number of active parameters, training dataset size, and the count of experts involved.

All the reviewers are positive about the contributions of the papers, including: (1) the novelty of the scaling laws framework that predicts a model’s training loss based on three core elements: the overall training budget, the number of active parameters during training, and the number of experts in the model; (2) the experiments are sufficiently extensive and appropriate to evaluate several aspects of the proposed scaling laws; (3) the writing and presentation of the paper are good.  After the rebuttal, most of the concerns of the reviewers were addressed, and all the reviewers are happy with the current stage of the paper.

In my opinion, the contributions and originality of the proposed scaling laws are sufficient for acceptance at ICML. Therefore, I recommend accepting it in its current form. However, I encourage the authors to address the reviewers’ suggestions and integrate their feedback into the camera-ready version of their paper.
Decision: Accept (poster)