Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training

• Large-scale LLM training leaves little scope for hyperparameter tuning despite the importance of well-chosen settings.Source: §1 · p. 1
• Existing approaches extrapolate optimal batch size and learning rate from small-scale fits, but it remains unclear whether compute, loss, model size, or data best explains scaling.Source: §1 · p. 2
• The paper fits power laws from hundreds of µP-trained models spanning weight decay, batch size, model size, and dataset size across 20–1280 tokens per parameter.Source: §1 · p. 2
• The AdamW timescale τ = B/(ηλD) follows a power law in tokens per parameter, enabling prediction of optimal weight decay across N, D, and B.Source: §1 · p. 2
• Pareto analysis finds that small, overtrained models can provide faster training and greater parallelism through larger D and Bcrit.Source: §1 · p. 2
• Both optimal and critical batch sizes scale as power laws in D, while a new Bcrit estimator avoids assumptions about a particular learning-rate schedule or optimizer.Source: §1 · p. 2

2.2 Methods: The AdamW timescale for LLMs, τ, and its scaling

• τ = B/(ηλD) normalizes the AdamW timescale by total optimization steps, representing the fraction of past iterations incorporated into final weights.Source: §2.2 · p. 3
• The proposed scaling law assumes τopt decreases as a power law in TPP and computes λopt from the µP-adjusted learning rate and τopt(TPP).Source: §2.2 · p. 3

2.3 Experimental details

• The experiments use GPT2-like models with µP, AdamW, SlimPajama, a held-out 1.1B-token validation set, and a linear warmup-then-decay learning-rate schedule.Source: §2.3 · p. 3

2.4 Results: τ and scaling λ

• Optimal τ remains stable as B scales within the usable regime, so λopt scales linearly with B; beyond Bcrit, τ drifts as gradient-based returns diminish.Source: §2.4 · p. 4, Figure 2 · p. 4
• When B changes, tuning λ is preferred to tuning η because η encounters instability at high values, whereas λ tuning was strictly superior in 6 of 8 comparisons.Source: §2.4 · p. 4
• Across model and dataset scales, τopt follows a precise decreasing power law, with R2=0.975 and values near 1.0 at 1 TPP and 0.01 at 1000 TPP.Source: §2.4 · p. 5
• The τ law generalizes across at least 3 orders of magnitude in compute, including held-out larger-scale points.Source: §2.4 · p. 5

3.1 Background: Bopt and Bcrit

• Bopt is the batch size minimizing loss at fixed N and D, whereas Bcrit is derived from the fitted trade-off curve as Dmin/Smin.B can be selected toward lower compute with smaller batches or toward lower training time with larger batches.Source: §3.1 · pp. 5–6
• Increasing B reduces optimization steps but increases tokens required for a target loss; Bcrit marks the transition where larger batches yield sharply diminishing step reductions.At B=Bcrit, reaching the target requires 2×Dmin tokens and 2×Smin steps; beyond Bcrit, much more data buys only small step reductions.Source: Figure 4 · p. 6, §3.1 · p. 6

3.2 Methods: estimating Bopt and Bcrit, and their scaling

• The method fits B-specific loss power laws, infers tokens needed for a target loss, and fits the resulting ⟨D,S⟩ pairs to obtain Bcrit.It does not assume a particular learning-rate schedule or optimizer and can estimate Bcrit at arbitrary losses without retraining.Source: §3.2 · p. 7

3.3 Results: Bopt and Bcrit

• Bopt and Bcrit obey power laws in D and Dmin, not in C or L.Fits achieve R2=0.984 for Bopt and R2=0.940 for Bcrit; fitted exponents are approximately 0.38 and 0.51, respectively.Source: §3.3 · p. 8
• The proposed D-scaling is fundamental: apparent dependence on TPP or model size arises when Bopt and Bcrit are plotted against C or loss instead of D.At fixed N or TPP, alternative plots can look approximately linear, but they do not fit all points across settings.Source: Figure 5 · p. 7, §3.3 · p. 8
• Projected DeepSeek batch sizes generally fall between Bopt and Bcrit, spanning the regime from greater compute efficiency to greater time efficiency.Because the DeepSeek law scales in C, its recommended B becomes larger for larger N.Source: Figure 5 · p. 7, §3.3 · p. 8
• Fitting batch-size laws with fixed λ distorts both Bopt and Bcrit estimates, while tuning λ to preserve the AdamW timescale yields cleaner scaling.Increasing fixed λ systematically raises estimated Bopt; the resulting slopes and R2 values generalize poorly to large-scale training.Source: §3.3 · p. 8
• Bopt ∝D^0.4 and Bcrit ∝D^0.5, enabling estimates at larger N and D and explicit FLOPs–training-time trade-offs across B.The proposed scaling relationships are intended to be estimated from small-N experiments with varied B, D, and tuned λ.Source: §3 · p. 5, §3.3 · p. 8

4.1 Background: compute-optimal and overtrained models

• The background frames the problem around allocating N and D at fixed compute, while noting that overtrained models can support greater data parallelism because Bopt and Bcrit rise with D.Compute-optimal training is associated with roughly 20 TPP, whereas overtrained models use higher D/N ratios for inference efficiency.Source: §4.1 · pp. 8–9

4.2 Methods: exploring the trade-offs of FLOPs vs. time

• The method constructs each frontier by fitting an iso-loss relation, varying B with the Bcrit-dependent excess-FLOPs model, and retaining non-dominated time–compute points.For a base setting reaching target loss with Dmin tokens, increasing B reduces time but requires proportionally more data and FLOPs through B/Bcrit(Dmin).Source: §4.2 · p. 9
• At fixed target loss and FLOP budget, time minimization balances compute-efficient base settings against overtraining, because higher D increases Bcrit while lower base FLOPs leave more budget for increasing B.The fitted iso-loss relation gives an optimal TPP ratio of approximately 20.6 for the reported loss target.Source: §4.2 · p. 9, §4.3 · p. 10

4.3 Results: balancing time and compute

• Overtrained models, but not undertrained models, lie on the FLOPs-versus-time Pareto frontier under Training Time ∝ Total FLOPs/B.Compute-efficient 20 TPP models remain optimal when minimizing FLOPs alone, while overtrained models can dominate undertrained models in both time and FLOPs.Source: §4.3 · p. 10
• Large-batch training should start from a 20 TPP+ base setting rather than target 20 actual TPP, because increasing B requires extra data to reach the same loss.A higher-B point trained from a 2 TPP+ base can already reach 20 actual TPP after a 10× FLOPs increase, making undertrained starting points Pareto-inefficient.Source: §4.3 · p. 10
• Bcrit estimation is essential because errors in its fitted scaling law shift the predicted frontier and can produce misleading Pareto-optimal configurations.Underestimating Bcrit can leave only low-TPP models appearing optimal, while increasing the exponent moves higher-TPP models onto the frontier.Source: §4.3 · p. 11
• The Pareto-optimal settings depend on the time and parallelism formulation: using time ∝ FLOPs/B favors higher-D overtraining, whereas time = steps produces a different frontier.The alternative formulation can make undertrained large models faster if larger models incur no time cost, although they use more FLOPs.Source: §4.3 · pp. 10–11, Figure 6 · p. 10

5 Conclusion

• Optimal τ decreases as a power law with tokens-per-parameter ratio, enabling systematic weight-decay selection across training scenarios.Source: §5 · p. 11
• Bopt and Bcrit both exhibit power-law dependence on dataset size D rather than primarily on compute, loss, or model size.Source: §5 · p. 11
• Smaller, overtrained models can be Pareto-advantageous when rapid training and high parallelism are prioritized.Source: §5 · p. 11
• Future analyses should include inference time and compute, model depth, and context length as additional Pareto objectives or configuration dimensions.Source: §5 · p. 11

A Broader impacts

• The methods can reduce training compute subject to time constraints, potentially lowering environmental and financial costs.Source: A Broader impacts · p. 17
• Scaling hyperparameters may reduce the financial burden of tuning and improve equity in AI research.Source: A Broader impacts · p. 17

Optimization and training setup

• The study uses AdamW and may not directly generalize to approximate second-order optimizers such as Shampoo.Source: B Limitations · p. 17
• The experiments use a single learning-rate schedule and do not verify whether jointly sweeping η and λ would yield substantially better models.Source: B Limitations · p. 17
• The results focus on single-epoch pre-training, while multi-epoch training may produce different optimal τ because of data repetition.Source: B Limitations · p. 17
• Only one dataset, vocabulary, and context length were tested, limiting conclusions about coefficient differences across setups.Source: B Limitations · p. 17

Precision

• Lower precision may leave Bopt and Bcrit scaling unchanged while altering τopt through a higher effective tokens-per-parameter ratio.Source: B Limitations · p. 17

Small batches

• Very small batches eventually degrade loss, and further tuning of Adam β parameters may be needed to improve them.Source: B Limitations · p. 18

Large batches

• Large-batch scaling laws omit bandwidth, communication, memory, synchronization, and optimizer-sharding effects that require system-specific profiling.Source: B Limitations · p. 18

C Experimental Details

• Around 400 LLMs were trained across the main experiments and supporting studies.Source: C Experimental Details · pp. 18–19
• Main experiments use SlimPajama, a 50,257-token GPT-2 vocabulary, 2,048-token context, AdamW, and cross-entropy validation loss.Source: C Experimental Details · p. 18
• All models use a 10% linear learning-rate warmup for a given tokens-per-parameter setting.Source: C Experimental Details · p. 18
• Experiments ran on Cerebras CS-3 systems; a 610M-parameter 20TPP model takes roughly six hours on one system.Source: C Experimental Details · p. 18
• Validation-loss variability was below 0.003 across repeated 111M-parameter, 20TPP runs.Source: C Experimental Details · pp. 18–19
• Table 3 records dataset sizes, tested batch sizes, and λ-tuning coverage, with untuned λ values inferred from the τopt law.Source: Table 3 · p. 19
• A 39M proxy model was used to tune 350 configurations of learning rates, initialization, and embedding/output scaling factors for µP.Source: C Experimental Details · p. 19

D.1 Optimizers for large-batch training

• A fitted Bcrit of about 12M tokens places BERT’s 4M-token LAMB run within the efficient range, while diminishing returns appear near 16.8M tokens.Source: §D.1 · p. 20
• Scaling λ rather than η with B supports stable training, but gradient redundancy still imposes an inherent limit on useful batch sizes.Source: §D.1 · p. 20

D.2 Critical batch size

• Prior definitions of Bcrit differ, including instantaneous loss-degradation thresholds and learning-rate scaling under constant schedules without weight decay.Source: §D.2 · p. 20
• The paper follows McCandlish et al. (2018)’s definition of Bcrit and measures it directly rather than relying only on gradient-noise-scale estimates.Source: §D.2 · p. 20
• The gradient noise scale provides only order-of-magnitude accuracy and can underestimate critical batch size.Source: §D.2 · p. 20

D.3 Detailed comparison with Zhang et al. (2024)

• The proposed Bcrit procedure fits B-specific loss power laws through a few validation losses, avoiding Zhang et al.’s need for frequent holdout evaluation.This approach is designed for learning-rate decay schedules and may reduce continuous-validation costs.Source: §D.3 · pp. 20–21
• The paper estimates Bcrit across many N and D settings, fits laws across model scales and TPP values, and evaluates fit quality with R2 and bootstrap parameter quantiles.Zhang et al. instead used three targeted experiments and fit their law using the 302M-parameter model.Source: §D.3 · p. 21
• Bcrit scales as D^0.462 in this work versus D^0.47 for Zhang et al. (2024), while their coefficient is 19% larger, suggesting similar underlying behavior.The paper reports its coefficient as 0.0471 and Zhang et al.’s converted law as Bzhang = 0.0559D^0.47.Source: §D.3 · p. 21

D.5 τ and effective learning rates

• Effective learning rates ηλ and τ are equivalent after batch normalization for fixed D, but tuning λ can achieve higher effective rates without the destabilizing effects associated with increasing η.Higher η can cause large early updates and larger weight norms, whereas the paper reports better results from adjusting λ with B.Source: §D.5 · p. 22

E.1 λ scaling with B

• For B < Bcrit, optimal λ increases linearly with B before plateauing or decreasing, so λ=0.1 is optimal only at particular B and TPP values.The figure caption independently summarizes the linear increase over small batch sizes until B exceeds Bcrit.Source: §E.1 · p. 22, Figure 7 · p. 23

E.2 Additional details on τ fitting

• The τ-fitting procedure minimizes a log-space parabola over validation-loss points, retaining the lowest loss at duplicate τ values and measuring B in tokens.Reported Bopt and Bcrit laws instead use sequences of 2048 tokens.Source: §E.2 · p. 22

E.3.1 Relationship to ηopt scaling laws in dataset size, D

• Eq. (11) converts the τ power law into an ηopt power law in D when λ, B, and N are fixed, matching the functional form used by Shen et al. (2024).The derivation gives mηD = −(mτ + 1) and cηD = N^mτ λ·cτ.Source: §E.3.1 · p. 24
• The fitted exponent mηD is −0.48 here versus −0.51 in Shen et al. (2024), while coefficient comparison requires converting sequence units and distinguishing adjusted from base µP learning rates.The paper converts Shen et al.’s coefficient to 0.0011 after dividing by their 4096-token sequence length.Source: §E.3.1 · p. 24
• The paper finds Shen et al.’s scaling law valid mainly for small models or width-only scaling, whereas its η, λ, and B laws jointly cover broader settings.It reports superior results from adjusting λ rather than η when optimizing τ.Source: §E.3.1 · pp. 24–25

E.3.2 Relationship to ηopt scaling laws in model size, N

• The paper recommends using µP to set ηopt at a model scale, then using the τ law to set λ according to batch size or dataset size.It states that either theoretical µP or empirical η laws can provide the model-scale learning rate.Source: §E.3.2 · pp. 25–26
• µP prescribes ηopt ∝ 1/W, which becomes ηopt ∝ N^-1/3 at fixed width-to-depth ratio and matches Porian et al. (2024)’s empirical law.The conversion uses N ∝ L W^2, fixed W/L, and therefore W ∝ N^1/3.Source: §E.3.2 · p. 25
• At constant τ, changing the number of steps or batch size can preserve the shape of weight-update contributions when D is fixed and B remains below the critical regime.The figure describes increasing B by K× while decreasing steps by K×, with λ adjusted to maintain τ.Source: Figure 8 · p. 26

E.4 The EMA perspective and learning rate schedules

• With a linear decay-to-zero schedule, matching τ produces the same normalized coefficient shape and therefore the same contribution pattern across training data at different step counts.The experiment compares 557, 5,568, and 55,680 steps while adjusting λ to preserve selected τ values.Source: §E.4 · pp. 26–27
• τ is independent of step count and provides a batch-independent AdamW timescale, but this equivalence breaks down when B > Bcrit because updates no longer contain linearly B-times more information.The EMA view identifies ηtλ as the time-varying smoothing parameter and treats model parameters as the EMA over weight updates.Source: §E.4 · pp. 26–27

F.1 Derivation of “extra data” Eq. (6)

• The Bcrit estimator derives D from the extra-data trade-off, computes Bcrit = Dmin/Smin, and fits a power law over the resulting Dmin–Bcrit pairs.The procedure fits B-specific loss laws first, then fits the Bcrit scaling law across model scales and target losses.Source: §F.1 · p. 27, §F.3 · p. 30

F.2 Estimating Bcrit

• Reliable Bcrit estimates require an irreducible-loss term, at least three points per power law, logarithmically sampled B values, and interpolation between fitting points.Beyond the fitting points, the curves may not predict behavior well, so the analysis restricts loss targets to interpolated regions.Source: §F.2 · p. 27

F.3 Estimating Bcrit for the 3.3B model

• For the 3.3B model, two B–Dmin pairs yield an estimated Bcrit of 4610 and Dmin of approximately 16 TPP.The estimate uses B=2016 at 23 TPP and B=4032 at 30 TPP, with nearly equal observed losses.Source: §F.3 · pp. 28–29

F.4 Bcrit scaling in loss

• At fixed TPP, assuming comparable N- and D-scaling exponents, loss becomes a power law in D; combined with Bcrit’s D power law, this yields a Bcrit power law in loss.The paper notes that Bcrit is not generally a power law in loss, but appears so when TPP is held fixed.Source: §F.4 · p. 30

• Larger fixed λ systematically increases estimated Bopt and lowers fit quality, while changing fitted parameters and projected batch sizes.Table 5 attributes these changes to conditional optima produced by fixed weight decay settings.Source: §F.5 · p. 30, Table 5 · p. 31
• Fixed λ makes Bopt|λ a conditional optimum that compensates for suboptimal τ rather than the globally tuned optimum.The resulting power law must balance maintaining a good AdamW timescale against staying below Bcrit.Source: §F.5 · p. 31
• For constant η, λ, and N, preserving τopt predicts B ∝ D^(m+1), with m + 1≈0.47, but fixed-λ fits deviate because gradient redundancy distorts small-D behavior.These distortions vary across D, altering the fitted slope and reducing generalization to larger-scale training.Source: §F.5 · p. 31
• Tuning λ removes the confound, producing a Bopt law consistent with D^0.4 scaling.Joint tuning prevents fixed weight decay from forcing batch size to compensate for timescale drift.Source: §F.5 · p. 31
• At 111M scale and target loss 3.03, fixed λ=0.1 lowers estimated Bcrit from 867 to 707, whereas tuned λ isolates gradient-redundancy effects.Across scales, fixed weight decay artificially steepens the Bcrit dependence; tuning yields a stable D^0.5 relation that generalizes.Source: §F.5 · p. 31