PaperlayerView exact revision on arXiv

Source-linked AI summary

Does Step Law Transfer to Small-Scale Language Models? An Empirical Recalibration Below 59M Parameters

Egor Romanyukov, Timofey Novikov, Timur Shokarov, Elizaveta Zorkina, Anastasia Palienko, Stepan Dergachev

arXiv:2609.27581v1cs.LGcs.CL

TL;DR

Step Law predicts optimal learning rate and batch size from model size and data, but it was not empirically tested below 59M parameters, a range relevant to small-model research and training. This paper tests transfer using loss-surface optima and finds that the power-law form holds with recalibrated coefficients, while direct transfer overestimates the optimal learning rate. The authors recommend using the recalibrated law or a roughly fourfold learning-rate adjustment within their studied regime.

  • Problem

    Step Law’s joint hyperparameter laws were not tested below 59M parameters, leaving their transfer to small models—used in interpretability and single-GPU training—unestablished.

    Source: §1 · p. 1, §2 · p. 3

  • Method

    The authors test Step Law on a small-model grid by extracting empirical optima from loss surfaces and comparing the original coefficients, recalibrated power laws, and a no-power-law outcome.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    The H2 outcome holds: the power-law form is preserved with different coefficients, and direct Step Law transfer overestimates the optimal learning rate by a median factor of about four.

    Source: §5.5.1 · p. 7, §5.5.2 · p. 7, §6 · p. 10

  • Takeaways & Limitations

    For the nanoGPT/TinyStories/BPE-2048 regime, use the recalibrated learning-rate law or roughly divide Step Law’s learning-rate prediction by four; B* remains primarily data-dependent.

    Source: §6 · p. 10

    The main fitted law covers N < 2.03M; the N = 5M cells are diagnostic anchors, so conclusions are restricted to the studied range.

    Source: §5.7 · p. 9

Abstract

from arXiv · show

Step Law gives power-law formulas for the optimal peak learning rate eta* and batch size B* when pre-training language models. It was calibrated on models between 59M and 1B parameters; the small-model regime N < 59M was never tested empirically by its authors. This regime matters for single-GPU training, interpretability research, educational experiments, and settings where larger models are infeasible on memory or cost grounds. We test whether Step Law transfers to small language models. We consider three outcomes: H1, the original coefficients work directly; H2, the power-law form holds but with different coefficients; and H3, a power law does not describe the optima in this regime. All experiments use a single nanoGPT/TinyStories pipeline with a 2048-token BPE vocabulary, AdamW, and a warmup-cosine schedule. The optimum for each (N, D) cell is extracted from the loss surface L(eta, B) via a local quadratic approximation in log-log coordinates over the smoothed training loss. The final dataset contains 29 unique (N, D) cells and 935 analysis-ready runs. The main refit uses 25 cells (815 runs) in the working range 4 <= D/N <= 600. On the pooled data we accept H2: the functional form is preserved, but the coefficients differ from the original. We obtain eta*(N, D) = 0.0985 N^(-0.508) D^(0.238) (R^2 = 0.834) and B*(D) = 3.6 x 10^(-4) D^(0.931) (R^2 = 0.950). Step Law's structural claim that B* is independent of N is reproduced (p = 0.87), but the growth of B* with D is nearly twice as steep as in the original work. Direct transfer of Step Law systematically overestimates the optimal learning rate: the median ratio eta_SL / eta* is approximately 4.0x, with a range of 2.4x to 6.6x.

1 Introduction

Step Law offers a way to predict optimization hyperparameters from smaller-scale experiments, but its calibration begins at 59M parameters. This paper tests whether its formulas transfer below that range, where small models matter for practical and scientific work.

1 Introduction

  • Exhaustive hyperparameter search at full scale is impractical because poor learning-rate or batch-size choices can waste millions of dollars in GPU time.Scaling laws aim to predict model behavior from substantially smaller-scale experiments.Source: §1 · p. 1
  • Step Law uses power laws to predict optimal learning rate and batch size, but its seven-model-size calibration covered only 59M–1B parameters.The original study left generalization beyond that calibrated range for future work.Source: §1 · p. 1
  • Small models matter for edge inference, interpretability research, and single-GPU training, including educational and early-stage projects.Models below 10M parameters trained on narrow, high-quality corpora can produce coherent text and rudimentary reasoning.Source: §1 · p. 2
  • Power laws can break down beyond their calibrated range, and reduced model capacity may change the hyperparameter landscape, making empirical testing necessary.The authors note that only an experiment can determine whether Step Law applies at small scale.Source: §1 · p. 2
  • The authors test three outcomes below 59M parameters: direct transfer of Step Law’s coefficients, a power law with recalibrated coefficients, or no valid power law.They construct loss surfaces over learning rate and batch size for each (N, D) pair using a single nanoGPT/TinyStories/BPE-2048 protocol.Source: §1 · p. 2

2 Related Work

Prior work describes loss scaling, optimization-hyperparameter scaling, and transfer across model widths, but the joint behavior of learning rate and batch size below 59M parameters remains unexplored on a standard parameterization. This study addresses that gap.

2 Related Work

  • Kaplan and Chinchilla relate language-model loss to scale and data, but do not specify optimal learning rate or batch size for a given (N, D).Their power-law scaling relationships can deviate when extrapolated far outside their calibrated range.Source: §2 · p. 2
  • AdamW and warmup-cosine scheduling must be fixed for meaningful comparisons because optimizer and schedule choices affect which learning rate is optimal.Warmup duration and ηmin can affect the loss landscape.Source: §2 · p. 2
  • Gradient-noise-scale reasoning links critical batch size to training progress, but its assumptions hold poorly in practice, leaving small-scale B* uncertain.Probe runs at different batch sizes have been proposed as an alternative that does not assume a particular loss geometry.Source: §2 · p. 2
  • DeepSeek-AI models optimization hyperparameters as functions of compute, while Step Law separates N and D and reports that B* depends only on D.Step Law was tested at 59M–1B parameters, leaving smaller models unvalidated.Source: §2 · pp. 2–3
  • µP predicts approximate learning-rate invariance with width, but requires architectural modifications and does not model data-volume dependence.The paper treats µP and Step Law as complementary: µP addresses width at fixed D, whereas Step Law also supplies D dependence.Source: §2 · p. 3
  • Small language models are used in interpretability, edge inference, and single-GPU training, where reliable guidance on η and B is needed.Prior work associates small models with coherent text, rudimentary reasoning, and mechanistic interpretability phenomena at the 1–10M scale.Source: §2 · p. 3
  • The gap is joint modeling of learning rate and batch size below 59M parameters on a standard parameterization.The comparison with prior work distinguishes the scale covered and whether loss or hyperparameters are modeled.Source: §2 · p. 3, Table 1 · p. 3

3 Problem Setup and Hypotheses

The paper defines empirical optima over learning rate and batch size for each model-size and data-budget pair, then tests whether Step Law transfers through three mutually exclusive outcomes. Its formal setup evaluates the power-law form and the original coefficients in the previously untested small-model regime.

3.1 Notation

  • A run fixes N, D, η, B, and seed; a sweep searches a predefined (η, B) grid for one (N, D) pair.Runs are treated as experimental observations mapped to smoothed training loss.Source: §3.1 · pp. 3–4

3.2 Formal Statement

  • The proposed family is η*(N, D) = c · N^α · D^β and B*(D) = d · D^γ, with batch size depending on D rather than N.The empirical optimum for each grid cell is found in the regime N < 59M; the research question asks whether the form and coefficients transfer and what mismatch costs in excess loss.Source: §3.2 · p. 4

3.3 Hypotheses

  • The hypotheses are H1, original coefficients hold; H2, the power-law form holds with changed coefficients; and H3, no power law fits acceptably at small scale.The three outcomes are pairwise exclusive and jointly exhaustive for the stated test.Source: §3.3 · p. 4

4 Background: Step Law

Step Law instantiates a power-law family for optimal peak learning rate and batch size, calibrated on large models. Its setup fixes a warmup-cosine schedule and uses smoothed training loss to assess optima.

4 Background: Step Law

  • Step Law gives η*(N, D) = 1.79 · N^-0.713 · D^0.307 and B*(D) = 0.58 · D^0.571.The batch-size law excludes N dependence.Source: §4 · p. 5
  • The coefficients were fit using 1,912 runs across seven model sizes from 59M to 1B and five data budgets, with about 0.094% relative deviation from the optimum.The original regression also confirmed that B* does not depend on N.Source: §4 · p. 5
  • The schedule uses 2,000-step linear warmup followed by cosine decay to ηmin = 10^-5, with η denoting peak learning rate.The experiments retain the same schedule class because landscape shape is schedule-sensitive.Source: §4 · p. 5
  • Smoothed training loss is used as an optimality proxy because it is reported to be minimized at the same η and B as validation loss.The paper separately tests the structural claim that B* depends only on D, motivated by gradient-noise-scale reasoning.Source: §4 · p. 5

5 Experiments

The experiments test Step Law on a small-model grid while preserving the comparison of optimal learning rate and batch size. The main refit uses 25 cells and 815 runs, with protocol choices adapted to this scale.

5.1 Protocol and Final Grid

  • The main experiments use nanoGPT-style models from 255K to 2.03M parameters, TinyStories, and an in-domain 2,048-token BPE tokenizer.Two 5M-parameter anchor cells are diagnostic and excluded from the main refit.Source: §5.1 · p. 5
  • 25 cells and 815 runs form the main refit within 4 ≤ D/N ≤ 600, drawn from a final grid of 29 cells and 935 analysis-ready runs.The two 5M-parameter cells and cells with D/N > 600 are retained for diagnostics rather than the main refit.Source: §5.1 · p. 5
  • All runs use AdamW and a warmup–cosine schedule with warmup set to 5% of the training budget, but use only one seed.The final protocol uses seed 1337; intervals therefore reflect variation across cells rather than initialization noise.Source: §5.1 · p. 5, §5.2 · p. 6

5.2 Design and Departures from the Step Law Setup

  • The study preserves comparison of the optimum’s location in (η, B) coordinates, but adapts the experiment to models tens to hundreds of times smaller than Step Law’s range.The design uses a 5% warmup because training steps vary with D/B, unlike Step Law’s fixed 2,000-step warmup.Source: §5.2 · p. 6

5.3 Tokenizer Choice and Its Consequences

The study uses a small in-domain vocabulary to avoid an oversized vocabulary projection dominating tiny models, and compares optimum locations rather than absolute losses across tokenizers. It extracts optima from smoothed training-loss surfaces and fits power laws in log-log coordinates.

5.3 Tokenizer Choice and Its Consequences

  • A 2,048-token TinyStories BPE tokenizer avoids making vocabulary embeddings and output projections disproportionately large relative to the small transformer backbone.Absolute cross-entropy values are not compared with GPT-2 BPE results because tokens carry different amounts of information.Source: §5.3 · p. 6

5.4 Metric and Optimum Extraction

  • For each (N, D) cell, the study builds a loss surface over (η, B) and extracts the optimum from a local quadratic fit within L ≤ Lmin + 0.15.The fit uses the smoothed training loss at the end of each completed run.Source: §5.4 · p. 6
  • If the fitted Hessian is positive definite and its vertex lies inside the explored box, that vertex is the optimum; otherwise, the observed grid minimum is used.On the 12 first-iteration cells, this method reproduces previously recorded optima exactly.Source: §5.4 · p. 6
  • The study fits log-log regressions for η* as a function of N and D and for B* as a function of D, then tests whether adding N improves the batch-size model.A hierarchical F-test evaluates the coefficient of log N in the extended batch-size model.Source: §5.4 · pp. 6–7

5.5 Final Results

On the main refit, direct Step Law transfer overestimates the optimal learning rate, while recalibration preserves the power-law form with different coefficients. The batch-size optimum remains independent of N, and the main coefficient patterns persist across iterations.

5.5.1 Direct Transfer of Step Law: Systematically Overestimated Learning Rate

  • The Step Law learning-rate prediction is about 4.0× the empirical optimum across 25 main-refit cells, with ratios ranging from 2.4× to 6.6×.The overshoot decreases as N grows; Table 3 reports overestimation by main-refit scale, and Figure 1 illustrates the shifted predictions.Source: §5.5.1 · p. 7, Table 3 · p. 7, Figure 1 · p. 7

5.5.2 Recalibrated Coefficients: H2 Instead of H1

  • The final refit gives η*(N, D) = 0.0985 N^-0.508 D^0.238 (R^2 = 0.834) and B*(D) = 3.6 × 10^-4 D^0.931 (R^2 = 0.950).The N exponent for η* is shallower than Step Law’s, while the batch-size D exponent is steeper.Source: §5.5.2 · p. 7

5.5.3 Batch Size Depends on D, Not on N

  • Adding N to the batch-size model yields a coefficient statistically indistinguishable from zero (p = 0.87), reproducing Step Law’s claim that B* depends on D, not N.The exponent comparison visual reports estimates against Step Law’s coefficients; the batch-size D exponent is significantly steeper.Source: §5.5.3 · p. 7, Table 4 · p. 8, Figure 2 · p. 8

5.5.4 Stability Across Iterations

  • Across iterations, α stays around 0.51–0.55 and γ around 0.93–1.02, while B* remains independent of N.The most sensitive quantity is the learning-rate D exponent β, which ranges from 0.20 to 0.35 across subsamples.Source: §5.5.4 · p. 8, Table 5 · p. 8

5.6 Verdict on H1/H2/H3

The verdict is H2: small-scale training retains the power-law structure but requires recalibrated coefficients, while optimal batch size remains independent of model scale.

5.6 Verdict on H1/H2/H3

  • H2 is the verdict: small scale preserves the power-law structure, but the coefficients need recalibration.The authors recommend against using Step Law verbatim and suggest taking about one quarter of its prescribed learning rate.Source: §5.6 · p. 8
  • Optimal batch size grows with data budget D as a power law, with no systematic separation by model scale N.Figure 3 uses point color to denote N.Source: Figure 3 · p. 9
  • Table 6 presents the pooled-data verdict on the three hypotheses.Source: Table 6 · p. 9

5.7 Limitations

The limitations are a single seed and a fitted range restricted to N < 2.03M; the hardware subsection compares the cost and speed of identical runs across GPUs.

5.7 Limitations

  • A single seed means the bootstrap confidence intervals reflect variation across cells, not full uncertainty across seeds.The learning-rate D-exponent β is especially sensitive and statistically compatible with the original value, while sitting at the margin of significance.Source: §5.7 · p. 9
  • The main fitted law is limited to N < 2.03M; the two N = 5M cells are diagnostic anchors, not fit data.Source: §5.7 · p. 9

5.8 Compute Cost and Hardware Choice

  • The hardware comparison measures the cost of processing one million training tokens for identical TinyStories runs.The metric uses hourly GPU price, wall-clock runtime, and training-token count.Source: §5.8 · p. 10
  • The A10 is 4.74× cheaper per token than the T4, with essentially identical validation loss; heavier sweep cells therefore ran on the A10.The A10 was 7.97× faster at 1.68× the hourly price.Source: §5.8 · p. 10, Table 7 · p. 10

6 Conclusion

The pooled results support recalibrating Step Law for the studied small-model regime rather than rejecting its power-law form, and the project provides reproducibility materials.

6 Conclusion

  • Across 935 runs, Step Law overestimates the optimal learning rate by a median factor of about four, while the power-law form is preserved and H2 is supported.The fitted batch-size law depends on data budget D alone, with no significant contribution from N.Source: §6 · p. 10
  • For small models in the nanoGPT/TinyStories/BPE-2048 regime, use the recalibrated learning-rate law or divide Step Law’s prediction by about four.Optimal batch size remains primarily set by D rather than N, with a steeper D slope in this range.Source: §6 · p. 10
  • The project provides code, configurations, run logs, optimum tables, and final refit scripts as a reproducible benchmark.Source: §6 · p. 10

A Per-Cell Optima

Table 8 reports robust-optimum estimates for 29 unique (N, D) cells, including cells excluded from the main refit.

A Per-Cell Optima

  • Table 8 lists robust 2D-quadratic optima and Step Law predictions for all 29 collected (N, D) cells.It reports the optimum estimates, Step Law learning-rate prediction, and their overshoot ratio.Source: A Per-Cell Optima · p. 12, Table 8 · p. 12
  • The table includes diagnostic cells excluded from the main refit: D/N > 600 cells and N = 5M anchors.Source: A Per-Cell Optima · p. 12

Related papers

Loading 2609.27581v1…