[Paper Review] Scaling Laws for Fine-Grained Mixture of Experts

Jakub Krajewski, Jan Ludziejewski+10 more|arXiv (Cornell University)|Feb 12, 2024

Advanced Statistical Process MonitoringDecision Sciences3 citations

TL;DR

This paper introduces granularity as a new hyperparameter to optimize Mixture of Experts (MoE) models, enabling fine-grained control over expert size. By deriving scaling laws that incorporate model size, training tokens, and granularity, the authors demonstrate that MoE models consistently outperform dense Transformers in efficiency, with compute savings exceeding 40× at large budgets—contrary to prior claims that MoE efficiency diminishes at scale.

ABSTRACT

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.

Motivation & Objective

• To address the high computational cost of training large language models (LLMs) by optimizing Mixture of Experts (MoE) architectures.
• To challenge the assumption that MoE efficiency diminishes at scale, particularly under fixed training duration and expert size.
• To identify optimal hyperparameters—especially granularity—for MoE models under varying computational budgets.
• To provide a unified framework for compute-optimal training of MoE models using scaling laws that include variable training duration and expert granularity.

Proposed method

• Introduces granularity (G) as a new hyperparameter that controls the size of experts in MoE models, allowing for fine-grained adjustment of expert capacity.
• Derives new scaling laws that model the relationship between model size, number of training tokens, granularity, and training duration, using FLOPs as the primary cost metric.
• Employs a modified MoE architecture with shared experts and segmented experts to enable efficient training at varying granularities.
• Conducts extensive experiments across multiple model sizes and compute budgets to validate the scaling laws and identify optimal configurations.
• Uses wall-clock time and perplexity as metrics to evaluate training efficiency, showing that higher granularity reduces training time for equivalent performance.
• Performs ablation studies on expansion rate (E) and demonstrates that findings hold across different E values, including E=16 and E=64.

Experimental results

Research questions

• RQ1 How does adjusting granularity in MoE models affect their training efficiency and performance across different computational budgets?
• RQ2 Do MoE models maintain or improve their efficiency advantage over dense Transformers as model size and training budget increase?
• RQ3 Is the common practice of setting expert size equal to the feed-forward layer size optimal for MoE models?
• RQ4 How do variable training durations and scaling laws interact to determine compute-optimal MoE configurations?
• RQ5 Can fine-grained MoE models achieve superior efficiency compared to dense Transformers even at extreme compute budgets (e.g., 10^25 FLOPs)?

Key findings

• MoE models with optimal granularity consistently outperform dense Transformers in terms of compute efficiency, with savings exceeding 40× at a 10^25 FLOP budget.
• The standard practice of setting expert size equal to the feed-forward layer (i.e., G=1) is suboptimal across nearly all computational budgets.
• The efficiency gap between MoE and dense models widens as model size and training budget increase, contradicting prior claims of diminishing returns.
• Optimal training duration varies with model size and granularity, and incorporating variable training time into scaling laws is essential for compute-optimal configurations.
• Fine-grained MoE models achieve lower perplexity faster, demonstrating significant gains in wall-clock training time compared to standard MoE and dense baselines.
• Scaling laws that include granularity, model size, training tokens, and variable training duration enable accurate prediction of optimal hyperparameters for any given FLOP budget.