Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling
2510.14717v1
cs.LG, cs.AI, math.OC, stat.ML
2025-10-18
Авторы:
Alexandru Meterez, Depen Morwani, Jingfeng Wu, Costin-Andrei Oncescu, Cengiz Pehlevan, Sham Kakade
Abstract
Increasing the batch size during training -- a ''batch ramp'' -- is a
promising strategy to accelerate large language model pretraining. While for
SGD, doubling the batch size can be equivalent to halving the learning rate,
the optimal strategy for adaptive optimizers like Adam is less clear. As a
result, any batch-ramp scheduling, if used at all, is typically tuned
heuristically. This work develops a principled framework for batch-size
scheduling and introduces Seesaw: whenever a standard scheduler would halve the
learning rate, Seesaw instead multiplies it by $1/\sqrt{2}$ and doubles the
batch size, preserving loss dynamics while reducing serial steps.
Theoretically, we provide, to our knowledge, the first finite-sample proof of
equivalence between learning-rate decay and batch-size ramp-up for SGD on noisy
linear regression, and we extend this equivalence to normalized SGD, a
tractable proxy for Adam, under a variance-dominated regime observed in
practice. Empirically, on 150M/300M/600M-parameter models trained at Chinchilla
scale using a constant (critical) batch size, Seesaw matches cosine decay at
equal FLOPs while reducing wall-clock time by $\approx 36\%$, approaching the
theoretical limit implied by our analysis.