PaperlayerView exact revision on arXiv

Source-linked AI summary

Predictable Scale: Part I, Step Law -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining

Houyi Li, Wenzhen Zheng, Qiufeng Wang, Hanshan Zhang, Zili Wang, Shijie Xuyang, Yuantao Fan, Zhenyu Ding, Haoying Wang, Ning Ding, Shuigeng Zhou, Xiangyu Zhang, Daxin Jiang

arXiv:2503.04715v7cs.LGcs.AI

TL;DR

LLM scale makes learning-rate and batch-size selection critical, while traditional grid search is prohibitively expensive, motivating transfer methods that extrapolate settings from smaller to larger models. The paper conducts 3,700 full LLM training runs and derives empirical scaling laws for optimal learning rate and batch size from the observed hyperparameter landscape. The predicted hyperparameters achieve a test-set loss within 0.094% of the global optimum found by exhaustive search and generalize across dense and sparse architectures and diverse data recipes.

  • Problem

    LLM scale makes learning-rate and batch-size selection critical, while traditional grid search is prohibitively expensive, motivating transfer methods that extrapolate settings from smaller to larger models.

  • Method

    The paper conducts 3,700 full LLM training runs and derives empirical scaling laws for optimal learning rate and batch size from the observed hyperparameter landscape.

  • Results

    The predicted hyperparameters achieve a test-set loss within 0.094% of the global optimum found by exhaustive search and generalize across dense and sparse architectures and diverse data recipes.

  • Takeaways & Limitations

    Step Law offers a plug-and-play way to reduce hyperparameter-tuning effort for industrial-scale LLM training.

  • Takeaways & Limitations

    The findings are primarily data-driven, and their theoretical basis and generalizability beyond the empirically validated domain remain future-work questions.

Abstract

from arXiv · show

The impressive capabilities of Large Language Models (LLMs) across diverse tasks are now well established, yet their effective deployment necessitates careful hyperparameter optimization. Although existing methods have explored the influence of hyperparameters on model performance, a principled and generalizable framework across model architectures and data recipes remains absent. In this study, we conduct an unprecedented empirical investigation training over 3,700 LLMs from scratch across 100 trillion tokens, consuming nearly one million NVIDIA H800 GPU hours to establish a universal Scaling Law for hyperparameter optimization in LLM Pre-training, called Step Law. We empirically observe that, under fixed model size ($N$) and dataset size ($D$), the hyperparameter landscape exhibits convexity with a broad optimum, substantially reducing the complexity of hyperparameter search. Building on this insight, we formally define and empirically validate the Step Law: The optimal learning rate follows a power-law relationship with $N$ and $D$, while the optimal batch size is primarily influenced by $D$ and remains largely invariant to $N$.Notably, our estimated optima deviate from the global best performance found via exhaustive search by merely 0.094\% on the test set. To our best known, Step Law is the first that unifies different model shapes and structures, such as Mixture-of-Experts models and dense transformers, as well as establishes optimal hyperparameter scaling laws across diverse data recipes. We contribute a universal, plug-and-play optimal hyperparameter tool for the community, which is expected to advance efficient LLM training at scale. All experimental code, data and checkpoints are publicly available at https://github.com/step-law/steplaw

1 Introduction

The paper introduces Step Law, a universal empirical framework for selecting learning rates and batch sizes across LLM architectures, sparsity levels, data recipes, and dataset sizes. It combines convexity of the hyperparameter loss landscape with scaling relationships whose predictions nearly match exhaustive-search optima.

  • Motivation: Hyperparameter optimization is important because learning rates affect divergence and convergence speed, while batch sizes trade off throughput against generalization.Existing transfer methods seek to reduce the prohibitive cost of exhaustive grid search.
  • Step Law formulation: Step Law models optimal learning rate as jointly dependent on model parameters N and dataset size D, while optimal batch size primarily depends on D.The proposed batch-size relationship assumes independence from N, and this assumption is statistically validated in the appendix.
  • Step Law results: The predicted hyperparameters produce a test-set loss only 0.094% above the exhaustive-search global optimum while outperforming baseline methods.The comparison is demonstrated on a 1B model trained on 100B tokens, where Step Law nearly matches the global optimum.
  • Loss-landscape analysis: The study identifies convex relationships between training loss and both learning rate and batch size when model size and dataset size are fixed.This convexity provides a broad, structured optimum for hyperparameter selection rather than an arbitrary landscape.
  • Generalization: Step Law generalizes across dense and sparse Mixture-of-Experts models, diverse architectures, data recipes, and sparsity configurations.Across these settings, the predicted learning rate and batch size yield the lowest validation smooth loss in the reported comparisons.
  • Experimental scope: The empirical investigation comprises 3,700 full LLM training runs covering approximately 100 trillion tokens and nearly one million H800 GPU hours.The resulting checkpoints and hyperparameter configurations are publicly available.

2 Related Works

Prior hyperparameter-transfer methods include theory-driven rules tied mainly to model architecture and data-driven scaling laws tied to model and dataset size. The paper positions Step Law as a unified alternative that predicts both learning rate and batch size across broader training dimensions.

  • Motivation: Hyperparameter transfer extrapolates optimal settings from smaller to larger models and is important because learning rate and batch size strongly affect pre-training performance.This motivates scaling rules that reduce the cost of direct optimization at large scale.
  • Theory-driven methods: Theory-driven methods began with µP learning-rate rules for model width but require custom initialization or modified attention mechanisms.These approaches do not cover data recipe, sparsity, dataset-size variation, or batch-size prediction.
  • Data-driven methods: Data-driven methods express hyperparameters as functions of model size N and dataset size D, but existing batch-size rules require prior loss estimates, ignore some dependencies, or assume fixed compute budgets.Earlier work linked learning rate to N and later incorporated D into joint scaling laws.
  • Step Law: Step Law claims that N and D alone suffice to predict both learning rate and batch size across model shapes, sparsity levels, and data recipes.The paper presents this as unified hyperparameter guidance for LLM pre-training.

3 Methodology

The methodology defines hyperparameter optimality under fixed architecture, data distribution, model size, and dataset size, then empirically investigates loss landscapes and scaling behavior across diverse models and training settings. It uses controlled experiments to derive and evaluate learning-rate and batch-size guidance, including dense and MoE architectures.

  • 3.3 Hyperparameter Landscape: Convex loss relationships for learning rate and batch size, together with a stable near-optimal plateau, make hyperparameter selection tolerant to small deviations.These properties provide the empirical basis for scaling-law development and reduce the need for exhaustive search.
  • 3.3.2 Impact of Final Learning Rate Schedule: The study compares conventional lrmin = lrmax/10 with a fixed lrmin = 10^-5 schedule because larger final rates can bias optimal settings toward lower learning rates.The conventional schedule can hinder final-stage convergence when peak learning rates are relatively high.
  • 3.3.3 Evaluation Metric Consistency: Training and validation loss identify the same optimum at learning rate 1.95 × 10^-3 and batch size 393,216, supporting smooth training loss as an optimization proxy.The corresponding reported smooth training and validation losses are 2.279 and 2.038, respectively.
  • 3.4 Scaling Laws: Experiments across model sizes and dataset scales reveal power-law scaling, with optimal learning rate depending on both N and D while batch size primarily depends on D.The study spans seven model sizes from 60M to 1B parameters and five dataset scales from 2B to 100B tokens.
  • 3.5 Comparisons with Existing Approaches: At the 1B-model test point, Step Law finds solutions near the global optimum and achieves validation performance competitive with the best configurations.The method consistently achieves the lowest validation smooth loss in the reported large sparse MoE evaluation.
  • 3.5 Generalization: Step Law remains accurate across six model topologies, diverse data conditions, and MoE sparsity levels, indicating invariance beyond a single dense architecture.Across varied model shapes, optimal learning rate and batch size remain in a narrow range; across sparsity levels, prediction error stays within 0.5%.

4 Conclusions

The paper presents universal learning-rate and batch-size scaling laws, supported by loss-landscape convexity, to make LLM hyperparameter optimization more efficient and generalizable.

  • 4 Conclusions: Universal learning-rate and batch-size scaling laws, supported by loss-landscape convexity, provide a practical approach for navigating LLM pre-training hyperparameters.The approach is intended to move beyond expensive grid searches and limited transfer methods.

A.1 Notation

This section defines notation for model architecture, model and dataset scale, training objectives, compute, and optimal hyperparameters. It also specifies the Transformer dimensions and attention-head count used to describe architecture space.

  • N denotes non-embedding model parameters, ˆN denotes total parameters excluding embeddings but including the model head, and D denotes dataset size in tokens.
  • LR is learning rate, BS is token batch size, L is cross-entropy loss, and C is compute budget measured in FLOPs.
  • η(N, D) and B(N, D) represent the optimal peak learning rate and optimal token batch size for a model with parameter count N and dataset size D.
  • A defines model architecture space using layer count, feed-forward dimension, hidden dimension, and attention-head count.
  • D also denotes the training data distribution governed by the data-generating probability distribution.

A.2 Model Scale Dominates Optimal Hyperparameter Selection Over Computational Complexity

Control experiments indicate that parameter count is a more robust determinant of optimal hyperparameters than computational complexity or detailed architectural proportions.

  • A.2 Model Scale Dominates Optimal Hyperparameter Selection Over Computational Complexity: Heatmaps compare loss across learning rate and batch size, with darker colors indicating lower loss values.The constant-compute visualizations show slightly more variance in optimal hyperparameter regions.
  • A.2 Model Scale Dominates Optimal Hyperparameter Selection Over Computational Complexity: Across architectural configurations, optimal hyperparameter regions remain consistent within each constant-scale experiment.The experiments vary hidden dimension, feed-forward dimension, attention heads, and layers while holding either N or M fixed.
  • A.2 Model Scale Dominates Optimal Hyperparameter Selection Over Computational Complexity: 2.4294 to 2.4776 minimum loss across constant-N models, versus 2.4346 to 2.5089 under constant-M experiments, indicates greater stability at fixed parameter count.The constant-N group also shows more consistent optimal hyperparameter regions than the constant-M group.
  • A.2 Model Scale Dominates Optimal Hyperparameter Selection Over Computational Complexity: Parameter count N appears more influential than specific architectural choices in determining optimal hyperparameter settings.This finding motivates scaling laws formulated in relation to N.

A.3 Model Structural Parameters

The study systematically varies dense and Mixture-of-Experts structural parameters to support comparisons across model configurations and matched compute or data scales.

  • A.3 Model Structural Parameters: Dense configurations vary dmodel, dff, Nhead, and Nlayer across models spanning 2.15 × 10^8 to 1.07 × 10^9 total parameters.
  • A.3 Model Structural Parameters: MoE configurations vary the number of experts, per-expert hidden size, top-k routing, and active parameter count.These variations enable direct dense-versus-sparse comparisons under matched compute budgets and data scales.

A.4 Composition of Training Datasets

The study uses four training-data recipes to test whether Step Law remains valid under shifts toward code, mathematics, and bilingual content.

  • A.4 Composition of Training Datasets: Four recipes compare baseline, enhanced code-and-mathematics content, more code-and-mathematics content, and English–Chinese bilingual data.The recipes are explicitly distinguished by dataset weight percentages and intended capability focus.
  • A.4 Composition of Training Datasets: The bilingual recipe augments the original English-only dataset with Chinese data to test validity in multilingual settings.
  • A.4 Composition of Training Datasets: The code-focused recipes incorporate 32.36% and 57.05% Code-The-Stack data, creating increasingly code-dominant distributions.

A.5 Statistical Validation of Batch Size Scaling Relationships

The statistical validation tests whether optimal batch size depends on model size and dataset size by comparing log-transformed regression formulations.

  • A.5 Statistical Validation of Batch Size Scaling Relationships: The full regression models log B as a function of log N, log D, or both to test batch-size scaling under power-law relationships.
  • A.5 Statistical Validation of Batch Size Scaling Relationships: OLS fitting and hierarchical F-tests assess whether model size and dataset size contribute to predicting log B.
  • A.5 Statistical Validation of Batch Size Scaling Relationships: The D-only formulation nearly matches the full model, with R2 = 0.821 versus 0.823, while the N-only model has R2 < 0.The comparison supports dataset size as the useful predictor and model size as insufficient on its own.
  • A.5 Statistical Validation of Batch Size Scaling Relationships: In the full model, log D is highly significant at p < 0.001, whereas log N is not significant at p = 0.257.

A.6 Loss Landscape Convexity

Across dense Transformer and MoE experiments, validation-loss surfaces are examined over learning rate and batch size, revealing broadly convex, bowl-shaped optima under fixed model and dataset conditions.

  • A.6 Loss Landscape Convexity: The 18 dense Transformer models use log–log learning-rate and batch-size axes, with smooth shifts in optimal regions and bowl-shaped minima.
  • A.6 Loss Landscape Convexity: MoE experiments evaluate 16 configurations while varying total parameter count and sparsity, sweeping learning rate and batch size over the same logarithmic grid.
  • A.6 Loss Landscape Convexity: The empirical findings remain primarily data-driven, motivating future theoretical derivations of the observed power-law relationships beyond the validated domain.
Loading 2503.04715v7…