PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Laws for Fine-Grained Mixture of Experts

Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michał Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygóźdź, Piotr Sankowski, Marek Cygan, Sebastian Jaszczur

arXiv:2402.07871v1cs.LGcs.AIcs.CL

TL;DR

The paper asks whether MoE efficiency persists as model sizes and training budgets grow, given limitations in prior fixed-training analyses. It introduces granularity and joint scaling laws over model size, training tokens, and granularity, then uses them to select compute-optimal configurations. The results report substantial MoE compute advantages over dense Transformers and show that feed-forward-sized experts are almost never optimal.

  • Problem

    The paper examines whether MoE models remain more efficient as training budgets and model sizes increase, addressing prior analyses with fixed training assumptions or dataset sizes.

  • Method

    The paper introduces granularity and derives scaling laws incorporating training duration, parameter count, and granularity to calculate optimal MoE training hyperparameters.

  • Results

    With optimal settings, MoE models can outperform traditional Transformers at any computing budget; a 10^20-FLOP MoE matches a dense model using 20× more compute, with savings exceeding 40× beyond 10^25 FLOPs.

  • Takeaways & Limitations

    Expert size should not generally be fixed to the feed-forward layer, because granularity enables more efficient MoE configurations across compute budgets.

  • Takeaways & Limitations

    Increasing granularity raises computational and communication costs and increases memory usage, mainly because routing operations grow with G.

Abstract

from arXiv · show

Mixture of Experts (MoE) models have emerged as a primary solution for reducing the computational cost of Large Language Models. In this work, we analyze their scaling properties, incorporating an expanded range of variables. Specifically, we introduce a new hyperparameter, granularity, whose adjustment enables precise control over the size of the experts. Building on this, we establish scaling laws for fine-grained MoE, taking into account the number of training tokens, model size, and granularity. Leveraging these laws, we derive the optimal training configuration for a given computational budget. Our findings not only show that MoE models consistently outperform dense Transformers but also highlight that the efficiency gap between dense and MoE models widens as we scale up the model size and training budget. Furthermore, we demonstrate that the common practice of setting the size of experts in MoE to mirror the feed-forward layer is not optimal at almost any computational budget.

1 INTRODUCTION

The paper studies whether MoE efficiency persists as training budgets and model sizes grow, challenging prior conclusions based on fixed training assumptions. It introduces granularity and scaling laws to identify compute-efficient MoE configurations.

  • Motivation: MoE methods target the high computational cost of training large language models while retaining comparable effectiveness.Prior models such as Switch and Mixtral demonstrated lower computational costs with comparable effectiveness.
  • Research question: The paper asks whether MoE models remain attractive as language-model training budgets increase.This question matters because prior studies reported narrowing efficiency gaps or possible dense-model advantages at larger scales.
  • Research question: The analysis relaxes assumptions of fixed training duration and constant expert size in prior research.These assumptions constrain how model size, training tokens, and expert configuration can vary together.
  • Results: A compute-optimal MoE trained with 10^20 FLOPs matches a dense Transformer trained with 20× more compute, while savings exceed 40× beyond 10^25 FLOPs.The paper also reports that matching expert size to the feed-forward layer is almost never optimal.
  • Contributions: Granularity controls expert size, enabling selection of more efficient MoE configurations.The paper introduces granularity as a hyperparameter for determining the optimal size of experts.
  • Contributions: The paper derives MoE scaling laws over training duration, parameter count, and granularity to calculate optimal training hyperparameters.The joint scaling law incorporates the variables needed to optimize training configurations.
  • Results: With optimal settings, MoE models can outperform traditional Transformers at any computing budget.This conclusion contrasts with prior results that predicted MoE underperformance beyond a certain model size.

2 RELATED WORK

Related work covers MoE routing methods, scaling laws for neural models, and prior analyses of MoE efficiency. The paper distinguishes its contribution by jointly analyzing training hyperparameters and granularity.

  • Mixture-of-Experts: MoE research progressed from sparsely gated recurrent layers to sparse Transformer layers and diverse routing strategies.Later work explored single-expert routing, assignment-based balancing, hashing, expert choice, and differentiable routing.
  • Mixture-of-Experts: Recent approaches segment experts or vary memory-block size, which can be interpreted as modifying granularity.The paper positions its work as providing a broader comparison of training hyperparameters and principled selection criteria.
  • Scaling laws: Scaling laws relate model loss to variables such as parameter count, training samples, and computational budget.Prior work established power-law relationships for dense Transformers and extended them to different training settings and architectures.
  • Scaling laws: Prior MoE scaling work found routed models more efficient only up to a certain model size when dataset size was fixed.The paper challenges this conclusion by considering the number of training tokens alongside model size and other hyperparameters.

3 BACKGROUND

Dense Transformers use feed-forward layers as their main parameter and FLOP cost, while MoE replaces these layers with selectively routed experts. Existing MoE scaling analysis was limited by fixed dataset size, which can cause undertraining as models grow.

  • Feed-forward layers contain most Transformer parameters and require the largest FLOP budget, motivating their replacement in MoE models.
  • MoE replaces the feed-forward layer with multiple experts, increasing total parameters while keeping computational cost approximately constant through subset routing.
  • Standard MoE experts typically mirror the original feed-forward dimensions, so total parameters scale linearly with the number of experts.
  • Dense Transformer loss is modeled as L(N, D) = c + a N^α + b D^β, relating loss to model size and training-token count.
  • Prior MoE scaling results used a fixed dataset of 130B tokens, limiting their applicability to other dataset sizes and potentially causing undertraining during model scaling.

4 GRANULARITY

Granularity controls expert size independently of expansion rate by splitting experts along their hidden dimension. Increasing it routes each token to more smaller experts while keeping active parameters constant, potentially improving expert-token matching.

  • Granularity G is defined relative to expert hidden dimension dexpert and the feed-forward dimension dff, with G = 1 representing the standard expert size.
  • The paper allows expert hidden dimensions to differ from the standard feed-forward layer, parameterizing MoE with granularity G and expansion rate E.
  • Increasing G splits experts into smaller units and routes each token to G granular experts, preserving the number of active parameters.
  • Expansion rate E compares total MoE-layer parameters with the standard feed-forward layer or, equivalently, total parameters with active parameters.
  • For non-granular models with G = 1, expansion rate equals the number of experts.
  • At fixed expansion rate, higher granularity gives the model more flexibility in mapping datapoints to experts and is incorporated into the scaling laws.

5 SCALING LAWS

The paper derives and fits scaling laws that model loss as a function of granularity, model size, and training tokens. Experiments support power-law behavior and show that compute-optimal comparisons require accounting for both training duration and routing cost.

  • 5.3 THE FORM OF THE JOINT SCALING LAW: The joint scaling-law objective incorporates granularity into a loss function over N, D, and G while preserving consistency with fixed-granularity laws.
  • Experiments: Over 100 experiments vary model sizes from 129M to 3.7B parameters, training durations from 16B to 130B tokens, and granularity from 1 to 16.
  • 5.1 POWER LAW WITH RESPECT TO GRANULARITY: For fixed N and D, increasing granularity lowers loss with approximately exponential returns converging to a positive constant.
  • 5.2 SCALING THE MODEL AND DATASET SIZE: At fixed granularity, loss follows a power-law relationship in model size N and training tokens D, analogous to dense-model scaling.
  • 5.4 FITTING THE PARAMETRIC SCALING LAW: The fitted scaling laws achieve RMSE = 0.015, and excluding the 20% lowest-perplexity models yields validation RMSE = 0.019.
  • 5.5 MOE SCALING PROPERTIES: Higher granularity lowers loss per training step, but excessive G can make routing cost dominate; compute-optimal selection therefore accounts for routing FLOPs.
  • 5.5 MOE SCALING PROPERTIES: MoE underfitting terms are 30.8D^-0.147 versus 26.7D^-0.127 for Transformers, indicating longer required training but better scaling after competitiveness is reached.

6 OPTIMAL ALLOCATION OF COMPUTATIONAL BUDGET

The paper allocates a fixed compute budget across model size, training duration, and granularity by modeling routing-aware FLOPs and numerically optimizing the resulting configuration. Higher granularity is favored at larger budgets, while compute-optimal MoE models remain more efficient than dense Transformers.

  • 6.2 OPTIMAL CONFIGURATION: The optimization selects N, D, and G for a fixed FLOPs budget F, subject to the routing-aware compute constraint.The resulting problem is reduced to one dimension and approximated with Brent’s method.
  • 6.1 COMPUTATIONAL COST OF GRANULARITY: Higher granularity lowers loss for the same training steps, but excessive granularity can make routing the wall-clock bottleneck.Increasing granularity also raises computational, communication, and memory costs.
  • 6.1 COMPUTATIONAL COST OF GRANULARITY: Routing-aware FLOPs account for both transformer computation and routing overhead, which increases with granularity.The routing term depends on d_model, expert expansion rate, granularity, and the number of training tokens and blocks.
  • 6.3 MOE IS ALWAYS MORE EFFICIENT: With optimal training hyperparameters, MoE models are more efficient than dense Transformers regardless of model size, and their efficiency gap increases with scale.Short fixed-duration schedules can instead leave large MoE models under-trained, explaining earlier contrary results.
  • 6.2 OPTIMAL CONFIGURATION: Compute-optimal training uses longer training for smaller active models, shorter training for larger ones, and higher granularity at larger compute budgets.These trends are reported for active parameter scales around 1B to 10B.

7 DISCUSSION

The discussion bounds the conclusions by identifying extreme-granularity failures, fixed expansion-rate experiments, hardware-dependent expansion costs, and setup-dependent granularity costs.

  • Extreme Granularity: At extreme granularity, performance can decline when routing parameters exceed the active parameters in the experts.The reported example is G = 64 with d_model = 256 and E = 64; such granularity is mainly useful for substantial models.
  • Varying Expansion Rate: The study primarily evaluates E = 64 because of computational constraints, although results for E = 16 preserve the main findings.Different expansion rates may be preferable for different target model sizes in memory.
  • Including E in the formula: A unified optimization over N, D, G, and E is left for future work because larger E improves performance but increases memory requirements.The preferred expansion rate may therefore depend on the available hardware configuration.
  • Modeling the cost of granularity: The precise cost of granularity depends on the training setup, hardware, and distributed-training implementation, especially through transfer costs.The paper models cost with FLOPs but also observes wall-clock gains for fine-grained MoE in its setup.

8 CONCLUSIONS

The paper concludes that granularity is a key MoE efficiency control and that compute-optimal scaling must vary both granularity and training duration. Its scaling laws support consistent MoE efficiency advantages over dense Transformers.

  • 8 CONCLUSIONS: Granularity G controls expert size, and the standard setting G = 1 is suboptimal across a broad range of FLOPs.The paper recommends higher granularity values to improve MoE performance and efficiency.
  • 8 CONCLUSIONS: The proposed scaling laws incorporate model size, training duration, and granularity to identify compute-optimal MoE configurations.The conclusion emphasizes varying training length alongside granularity.
  • 8 CONCLUSIONS: MoE models consistently outperform dense Transformers in efficiency and scaling under the paper’s optimized settings.The conclusion presents these findings as practical guidance for improving large-language-model computational efficiency.

A ARCHITECTURE AND TRAINING SETUP

The experiments use decoder-only Transformers with a shared language-modeling setup for dense and MoE models, while the MoE implementation uses balanced Expert Choice routing. Scaling-law coefficients are fitted on held-out validation data.

  • A ARCHITECTURE AND TRAINING SETUP: All models are decoder-only Transformers trained on C4 with GPT2 tokenization and AdamW optimization.Runs use 0.5M-token batches packed into 2048 sequences, a 2e−4 peak learning rate, and cosine decay to 2e−5.
  • A ARCHITECTURE AND TRAINING SETUP: MoE models use Expert Choice routing to balance expert load without tuning additional hyperparameters.Tokens are grouped by position across sequences, with group size fixed at 256.
  • A ARCHITECTURE AND TRAINING SETUP: The MoE and dense variants are matched in FLOPs for the same d_model.The implementation also applies router softmax over experts and adds layer normalization before the MoE output.
  • A ARCHITECTURE AND TRAINING SETUP: The study reports separate architecture and training variants for MoE and dense models.These variants are listed in Tables 3 and 4.
  • A ARCHITECTURE AND TRAINING SETUP: Scaling-law coefficients are fitted after excluding the 20% of datapoints with the lowest perplexity for validation.The fitted coefficient values are reported separately.

C RELIABILITY OF COMPUTE OPTIMAL FORMULA

The section evaluates the stability of the compute-optimal predictions using bootstrap-derived uncertainty intervals.

  • 10th and 90th percentiles are estimated by bootstrapping 80% samples of the data 100 times.These percentiles are used to assess prediction stability.
  • The resulting uncertainty details are reported in Table 6.
  • The bootstrap analysis follows the procedure used by Hoffmann et al. (2022).

D VARYING EXPANSION RATE

This section studies MoE models with expansion rate E = 16, fits the scaling-law coefficients, and derives compute-optimal training parameters while accounting for routing overhead.

  • Varying expansion rate: E = 16 models use the training procedure described in Appendix A and are listed in Table 7.
  • Varying expansion rate: The scaling-law coefficients are fitted with the procedure from Section 5.4 and reported in Table 8.
  • Varying expansion rate: Compute-optimal training parameters are derived from the fitted coefficients and FLOPs calculation formulas, with results in Table 9.
  • Varying expansion rate: Larger compute budgets imply larger optimal values of G, as observed for E = 64.For E = 16, the reported 10th and 90th percentiles form wider intervals because fewer experiments and shorter training durations were used.
  • FLOPs calculation: The Transformer-training FLOPs formula accounts for routing-operation overhead in MoE.The calculation uses cf = 6 for active-parameter and token pairs, while the routing constant is assumed to be cr = 14.
  • FLOPs calculation: The routing constant may range from 6 to 20 depending on implementation, but the main conclusions are resistant to different assumptions.FLOPs are counted in multiplication-addition pairs.
  • Scaling visualizations: Figure 6 illustrates scaling N and D at fixed granularity values G = 1, G = 2, G = 8, and G = 16.
  • Scaling visualizations: Figure 7 illustrates granularity scaling with fixed N and D across four configurations.The configurations use N = 64 × 25M or 64 × 49M, D = 16B or 32B, and constants from 2.88 to 3.12.
Loading 2402.07871v1…