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

Key takeaways

• The move is to scale a derived quantity rather than the hyperparameter. Under AdamW the weights are themselves an EMA of updates, with timescale τ = 1/(η λ) in steps. Normalized by the total step count S, the fraction τ̃ = 1/(η λ S) says what share of training is effectively averaged into the final weights. λ is then read off from τ̃ rather than set directly.
• τ̃ is not constant, and that is the correction. Wang and Aitchison proposed holding it fixed as scale changes. Across tokens-per-parameter (TPP) from 20 to 1280, the optimal τ̃ follows a power law with exponent ≈ −0.52 (bootstrapped 10th/90th percentiles −0.529 to −0.507, R² = 0.975) — roughly 1.0 at 1 TPP falling to 0.01 at 1000 TPP. Held-out points at up to 1000× the fitting compute sit on the line.
• With muP in place, tune λ rather than η. Against the usual practice of sweeping the learning rate at a default λ = 0.1, tuning λ at the muP-transferred learning rate was strictly better in 6 of 8 cases.
• Optimal batch size scales with the token budget, not with compute. B_opt ∝ D^0.38 (10th/90th 0.367–0.391, R² = 0.984). The apparent power law in compute that prior work fitted holds only at a fixed TPP; it is a shadow of the D dependence.
• Critical batch size likewise: B_crit ∝ D^0.47 (R² = 0.940), against Zhang et al.’s independently measured 0.462 — different context length, architecture, dataset, parameterization, schedule and tuning strategy, and the exponents agree to within about 2%.
• A method for B_crit that does not need a fixed training duration, which is what made the quantity awkward to measure under a decaying schedule.
• The Pareto analysis over training time against compute favours small, over-trained models: they take fewer wall-clock steps and admit more data parallelism, because B_crit grows with D.
• GPT-2-like architecture with ALiBi and SwiGLU, muP, SlimPajama, linear schedule with 10% warmup, models to 3.3B.

Standing in the anthology

It retires an exponent the record has carried since its first month. SOTA-097 says optimal batch size scales as B ∝ C^0.24, from Kaplan’s equation 1.7, and its own body already said the right thing: “that the batch exponent has not been re-derived by anyone since is a reason to hold it loosely”. It has now been re-derived twice, by groups who did not coordinate, and the scaling variable is not compute.

The paper it corrects is Kaplan et al., LIT-028, and only on that exponent: equation 1.7 fits optimal batch size as a power law in compute, and this study finds optimal and critical batch size both scale in the token budget D, with the compute law a projection that holds only at fixed tokens-per-parameter.

That is the rarer and more valuable kind of correction: not a better fit of the same relationship but a different independent variable. Kaplan’s fit is not wrong within its slice — at fixed TPP, D and C move together and either fits — it is a measurement of a projection presented as the thing itself.

It also supplies what SOTA-255 said nobody had. That practice records Kim et al.’s finding that optimal weight decay under data constraint is ~30× the inherited 0.1, and its conditions section says plainly: “the paper gives no way to predict the optimum from the parameter-to-token ratio, so the practice is really ‘sweep it, and sweep upward’.” Here is a law in exactly that ratio.

The two do not compose, and the record should not pretend they do. This study normalizes by S on the explicit basis that “LLM pre-training only uses one epoch of data”, and fits from 20 to 1280 TPP. Kim et al. epoch a fixed corpus many times at parameter-to-token ratios above Chinchilla — TPP below 20, and often below 1. The two laws are fitted in disjoint regimes under different assumptions about repetition, and nobody has checked whether either extrapolates into the other’s. What the record holds now is two answers to “how should weight decay move under data pressure”, each well-measured where it was measured.