Towards joint scaling laws with optimal batch size schedules

2607.27731 ↗ · 2026-07-30 · cs.LG, cs.AI, math.OC

Jiaxiang Li, Zhiqi Bu, Shiyun Xu

The finding

Headline claim

Replacing static batch sizes with a closed-form schedule derived from the learning rate schedule shape improves compute efficiency in LLM training by 6 to 15 percent at zero extra computational cost.

The test

Theory: convex-optimization loss-prediction bounds turned into a constrained optimization over batch-size sequences under data or compute budgets. Empirics: pre-training runs on Llama3 (dense) and Qwen3 (MoE) up to 7B parameters on Fineweb-edu under a fixed 1e20 FLOPs budget, plus fine-tuning of vision-language models and math-domain models, with AdamW and Muon optimizers and ablations across scales and configurations.

Effect size & scope

Authors report 6 to 15 percent compute-efficiency improvement over static batch size baselines, and lm-evaluation-harness scores favor dynamic schedules on Llama3 (e.g. LAMBADA 35.36 vs 32.23) and Qwen3 (LAMBADA 37.71 vs 36.06). They also argue the η/B ratio is not a sufficient statistic: matched η/B trajectories differ by +0.017 to +0.743 validation loss depending on which hyperparameter realizes the schedule. The claimed optimal batch size schedule depends only on learning rate schedule shape, not peak learning rate, model size, architecture, or optimizer. Scope: validated at up to 7B parameters and 1e20 FLOPs, far below frontier runs.

Limitations

The theory assumes convex loss and SGD, extended to non-convex networks and adaptive optimizers only empirically. The gradient-noise model assumes a time-independent critical batch size. Experiments are at most 7B parameters, and the paper does not demonstrate frontier-scale validation. The efficiency gain is measured on their own training configurations and benchmarks, and the full text provided omits wall-clock and systems-throughput implications of dynamic batch sizes, which complicate real data-parallel training pipelines.

Counterpoint

A skeptic would note the convex-dominance story is an empirical approximation asserted, not proven, for non-convex deep networks; benchmark deltas like LAMBADA +3.1 points could be seed- or checkpoint-sensitive with no error bars reported; the 6 to 15 percent efficiency figure depends on how compute is normalized; and dynamic batch sizes impose engineering costs (pipeline re-balancing, communication patterns) that the 'zero computational overhead' framing may understate. Meta affiliation raises the question of lab-internal tuning that favors their stack.

Economy implication

Little or nothing yet for the AI economy in the next 12 to 24 months: if the 6 to 15 percent training-compute efficiency gain holds at frontier scale, it would slightly lower the marginal cost of pre-training runs and marginally weaken 'more compute = more value' arguments, but the paper validates only up to 7B parameters, so no staffing, capex, or pricing conclusion is warranted yet.

Who this matters for

Technical Capability: Training engineers get a free, scale-free batch-size recipe that is decoupled from learning-rate tuning, though it remains unproven at frontier scale.

Market & Valuation: If dynamic batch-size schedules transfer to frontier runs, they shave a mid-single to low-double-digit percentage off pre-training compute cost, a small but real margin effect for labs and hyperscalers spending billions on training.

Bears on (advisory · papers never write evidence)

[6.2] Algorithmic / training-efficiency gains reduce the compute (and chip-access) needed for frontier capability.

[B3] Training-efficiency gains will reduce aggregate compute demand (threatening "more compute = more value").

Abstract

Modern deep learning typically keeps the batch size static throughout training, thus overlooking the joint effect of learning rate and batch size on the training dynamics. In this paper, we study the deep learning dynamics through the lens of convex optimization and derive a joint characterization of loss in terms of both schedules, applicable to general optimizers and model architectures. This characterization yields a closed-form optimal batch size schedule for any prescribed learning rate schedule, and further leads to joint scaling laws that consistently outperform static batch size baselines, highlighting the significance of dynamic batch size schedule in large language model training.