PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2502.05172v2cs.LGcs.AIcs.CL

TL;DR

The compute–memory trade-offs of MoE models remain unclear, despite their computational efficiency. This paper derives joint scaling laws for dense and MoE models and finds that MoE can be more memory-efficient, with a rule of thumb for selecting expert counts under fixed memory and compute budgets.

  • Problem

    The compute–memory trade-offs of MoE models remain unclear, including whether MoE can be optimal when memory is constrained.

    Source: §1 · p. 1

  • Method

    The paper derives joint scaling laws for dense and MoE models that account for active parameters, dataset size, and expert count.

    Source: §1 · p. 1

  • Results

    For a fixed total parameter count, an MoE with E ≤ 8 experts outperforms a compute-optimal dense model when trained on E times more tokens at the same memory footprint.

    Source: §4.5 · p. 7

  • Takeaways & Limitations

    The scaling laws provide a framework for selecting an optimal MoE configuration under compute and memory constraints.

    Source: §1 · p. 1

    FLOPs may not reflect wall-clock training time across architectures, and the analysis does not consider differing implementations or distributed training algorithms.

    Source: §6 · p. 10

Abstract

from arXiv · show

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. To derive and validate the theoretical predictions of our scaling laws, we conduct over 280 experiments with up to 2.7B active parameters and up to 5B total parameters. These results offer actionable insights for designing and deploying MoE models in practical large-scale training scenarios.

1. Introduction

The paper addresses unclear compute–memory trade-offs for MoE models with joint scaling laws and reports that MoE can outperform dense models under matched budgets. It also offers guidance for choosing expert counts under those constraints.

1. Introduction

  • The paper addresses uncertainty about compute–memory trade-offs as MoE models grow, when adding experts has relatively little FLOP cost but can sharply increase memory requirements.Source: §1 · p. 1
  • MoE can be more memory-efficient than dense models, with the same loss and training budget achievable at lower memory usage.Figure 1 validates the memory-optimality thesis with a compute-matched MoE comparison to an overtrained dense model with 1.1B total parameters.Source: §1 · p. 1, Figure 1 · p. 2
  • The joint scaling law accounts for active parameters, dataset size, and expert count to analyze performance under strict memory constraints.The paper reports over 280 models, scaled to 2.7B active and 5B total parameters.Source: §1 · p. 1
  • The optimal expert count, including dense models with E = 1, depends on compute and memory constraints, and the optimal token-to-parameter ratio depends on E.Source: §1 · p. 2
  • Under identical compute and total-memory budgets, the 1.1B-parameter MoE models achieve lower final loss than the dense model and higher inference performance.The paper reports that MoE can often be preferable even when GPU memory is the constraint.Source: §1 · p. 2

2. Related Work

Prior work studies scaling laws for dense and MoE models, but differs in which variables and efficiency constraints it considers. This paper positions its analysis around expert count, training data, inference cost, and total memory.

2. Related Work

  • In Transformer LLMs, the standard Switch MoE replaces the feed-forward MLP with expert MLPs and routes each token to one expert.The layer adds a differentiable load-balancing loss to encourage even token-to-expert assignment.Source: §2 · p. 2
  • Dense-model scaling laws relate final loss to model size and dataset size, with later formulations incorporating variable cosine-cycle lengths and adjusting the equation form.Source: §2 · p. 3
  • Scaling-law studies also address other modalities, machine translation, and pruning-based sparsity under fixed compute.Source: §2 · p. 3
  • Earlier MoE scaling work varied model size and expert count on a fixed dataset, finding routed models more efficient only up to a certain size.Source: §2 · p. 3
  • Later work considered variable dataset size and expert granularity, while this paper keeps experts non-granular and varies expert count and training tokens.Sardana et al. (2024) assumed a fixed joint inference and training budget; this paper additionally considers accelerator memory as a limit.Source: §2 · p. 3
  • Unlike concurrent sparsity-focused analysis, this work includes inference cost and total memory and derives a method for scaling learning rate with expert count and model size.Abnar et al. (2025) studied optimal sparsity and the relationship between pretraining loss and downstream performance.Source: §2 · p. 3

3. Joint MoE Scaling Laws

The joint law extends fixed-expert scaling by modeling loss as a function of active parameters, training tokens, and expert count. Its form incorporates interactions among these variables and is empirically reported to fit MoE performance.

3. Joint MoE Scaling Laws

  • For a fixed expert count, the starting form follows established scaling-law practice, and its irreducible-loss term is independent of architecture.The constant represents irreducible loss from the dataset’s inherent entropy.Source: §3 · p. 4
  • The law’s interactions reflect observations that expert scaling follows a power law, MoE benefits diminish with larger models at fixed tokens, and grow with more tokens at fixed model size.The cited observations are attributed to Clark et al. (2022) and Ludziejewski et al. (2024).Source: §3 · p. 4
  • The proposed interactions extend the earlier form, which is recovered when the second and third terms are ignored.Source: §3 · p. 4
  • The authors fit the law’s coefficients using their experiments and report a good empirical fit for modeling MoE performance.Training-run and fitting details are presented in Section 5.Source: §3 · p. 4
  • The expert-count term uses a transformation that accounts for overhead at small expert counts and diminishing returns at very large counts, following Clark et al. (2022).Source: §3 · p. 4
  • The final law models loss jointly with active parameters, training tokens, and a transformed expert count, including expert-dependent exponents for the parameter and dataset terms.The paper reports that these interactions accurately model MoE performance.Source: §3 · p. 4

4. Compute and Memory Optimality

The section extends scaling-law optimization to memory and inference constraints, showing that expert count should be chosen jointly with model size, training data, and budget.

4.1. Compute Optimality

  • At fixed compute, more experts shift the optimum toward fewer active parameters and more training tokens, while their compute savings over dense models grow with expert count.Source: §4.1 · p. 5

4.2. Model Memory Optimality

  • At fixed compute, more experts improve performance when model size and training tokens are adjusted; 1.1B MoE counterparts also beat a size- and FLOP-matched dense model after dense overtraining.Source: §4.2 · p. 6

4.3. Total Memory Optimality

  • Including the KV cache makes MoE models compare more favorably, and increasing cache length moves their advantage over dense models to smaller model sizes.Source: §4.3 · p. 6

4.4. Inference Optimality

  • With training and inference sharing a compute budget, MoE models outperform dense models at smaller scales than under compute-only optimization because they use fewer inference FLOPs.Source: §4.4 · p. 6, Figure 3 · p. 7

4.5. Summary

  • For fixed total parameters, an MoE with E ≤8 experts trained on E times more tokens outperforms a compute-optimal dense model at the same memory footprint.Source: §4.5 · p. 7
  • Optimal E varies with training budget and memory limit; Table 2 evaluates RTX4090, H100, and 8xH100 settings with a 16k-token KV cache.Source: Table 2 · p. 7, §4.5 · p. 7
  • The E ≤8 rule is conservative because highly overtrained-model loss predictions are uncertain, and large-dataset availability varies by domain.Source: §4.5 · p. 7

5. Fitting the Scaling Law

The experiments fit a joint scaling law using extensive model runs and training choices intended to yield robust predictions. The fitting procedure includes expert-dependent learning-rate scaling and tests the formula on held-out and extrapolation data.

5.1. Model Hyperparameters

  • The models are decoder-only Transformers trained on FineWeb-Edu, using Switch layers, one expert per token, SwiGLU, mixed precision, and dropless evaluation.The study also uses standard router z-loss and load-balancing loss values, a GPT-2 tokenizer, and reduced-scale truncated-normal initialization.Source: §5.1 · p. 8

5.1.1. BATCH SIZE RAMP-UP

  • Batch size ramps from 64K to 128K after 0.5B tokens, then to 256K after 1B tokens, using a grid to select the transition point.The ramp-up is intended to mitigate potential performance issues from excessively large batches, especially early in training.Source: §5.1.1 · p. 8

5.1.2. LEARNING RATE SCALING

  • Increasing E requires a lower learning rate; the fitted rule uses active parameter count and expert count and predicts nearly optimal rates at E = 4 and E = 32.The rule is fitted on E = 1 and E = 8, then validated at E = 4 and extrapolated to E = 32; omitting its E dependence can impair MoE performance, especially at E = 32.Source: §5.1.2 · p. 9

5.1.3. LEARNING RATE SCHEDULE

  • Training uses a constant learning rate with linear warmup over 130M tokens and linear decay to zero over the final 20% of tokens.Longer runs reuse intermediate checkpoints from shorter runs.Source: §5.1.3 · p. 9

5.2. Optimization of Formula Coefficients

  • 0.018 maximum absolute error on the held-out extrapolation set, with prediction errors across expert counts showing no bias toward a specific E.The joint fit has RMSEv = 0.0039 on a held-out set of 30 lowest-loss runs and RMSEt = 0.0062 on the training data; separate-E fits lower training error but are more prone to overfitting.Source: §5.2 · p. 9, Figure 5 · p. 10

6. Limitations and Future Work

The analysis focuses on standard MoE experts and uses FLOPs as a proxy for compute, while assuming Chinchilla scaling laws. The authors identify boundaries involving architecture, implementation, and extreme training regimes.

6. Limitations and Future Work

  • The study covers standard MoE, with each expert matching the dense model’s feed-forward layer size; fine-grained and dropless variants are left for future work.The authors suggest these variants may further improve the reported MoE benefits.Source: §6 · p. 10
  • FLOPs may not reflect wall-clock training time across architectures, and the analysis does not account for different implementations or distributed training algorithms.Comparing total rather than active parameters partially alleviates the issue because of the shared memory bottleneck.Source: §6 · p. 10
  • The formulas assume Chinchilla scaling laws, which have known limitations at extreme token-to-parameter ratios; fit quality also regressed for some heavily undertrained or overtrained runs.Source: §6 · p. 10

7. Conclusions

The joint MoE scaling laws relate loss to model parameters, training tokens, and expert count, and show that MoE models can outperform dense models under memory or total-parameter constraints. The analysis also offers a framework for choosing expert counts under given constraints.

7. Conclusions

  • MoE models can outperform dense models under memory or total-parameter constraints, contrary to the intuition that they are more memory-intensive.The scaling laws relate loss to parameters, training tokens, and expert count.Source: §7 · p. 10
  • Optimal training strategies shift with expert count, providing a framework for selecting MoE hyperparameters under specified constraints.The analysis highlights trade-offs between memory and compute performance.Source: §7 · p. 10

A.1. Counting Parameters

  • Parameter counts distinguish active from total parameters and whether embedding and unembedding layers are included; by default, Nact and Ntotal include those layers.The paper notes that Kaplan et al. (2020) count nonembedding parameters, while Hoffmann et al. (2022) include embedding and unembedding parameters.Source: §A.1 · p. 14

A.2. Counting FLOPs

The FLOPs accounting specifies training and inference costs in terms of active parameters and tokens, treats routing costs as negligible, and uses a parameter-based rule to scale model configuration.

A.2. Counting FLOPs

  • Training and inference FLOPs are modeled as Ftraining = 6NactDtraining and Finference = 2NactDinference, respectively.The assumptions are based on Sardana et al. (2024).Source: §A.2 · p. 14
  • With at most 32 experts and implicit expert granularity of 1.0, routing memory and FLOPs costs are considered negligible.The paper follows Clark et al. (2022) for this assumption.Source: §A.2 · p. 14

A.3. Model Configs

  • Most experiments scale configurations using Nblocks = Nheads = dmodel/64, and the same relations are assumed in calculations.The rule is based on Kaplan et al. (2020).Source: §A.3 · p. 14

B. Fit Details

The fit details describe coefficient optimization for the joint formula and its reduction to Chinchilla scaling laws at fixed expert counts, where the dataset exponent increases significantly.

B. Fit Details

  • Table 3 presents the fitted coefficients of the joined formula.The supplied passage identifies the table's scope but does not give coefficient values.Source: Table 3 · p. 14
  • The dataset exponent ν increases significantly when the joint formula is reduced to the Chinchilla scaling law for a given expert count, changing compute-optimal parameter-to-token ratios.Source: Table 4 · p. 14
  • The joint formula's coefficients are fitted with LBFGS using Huber loss, with underperforming or outlier runs downweighted proportionally to loss.The fit uses δ = 0.01 and selects coefficients minimizing the sum of training and held-out extrapolation RMSE.Source: B. Fit Details · pp. 14–15

C. Compute- & Memory-Matched Models

The learning-rate fit ablation supports including expert count when extrapolating to E = 32. Compute- and memory-matched comparisons are bounded by the scaling law's token-to-parameter coverage.

C. Compute- & Memory-Matched Models

  • 155M activated parameters trained on 64B tokens yield a token-to-parameter ratio of approximately 414 for an E = 16 MoE matched to a 1B dense model trained on 10B tokens in FLOPs and memory.This ratio exceeds the dataset-covered range, so higher E values are not considered under the constraint.Source: Figure 6 · p. 15

D. Learning Rate Scaling Fit

  • Figure 7 visualizes the learning-rate rule fit at E ∈{1, 8}, interpolation at E = 4, and extrapolation at E = 32.Source: Figure 7 · p. 16
  • The E-omitting fit extrapolates clearly suboptimally at E = 32, supporting the need to include expert count in the learning-rate scaling rule.Training-set performance for E ∈{1, 8} looks acceptable.Source: Figure 8 · p. 16

E. Experiments Listing

The experiment listing catalogs model configurations by architecture, active parameters, expert count, and dataset size.

E. Experiments Listing

Related papers

Loading 2502.05172v2…