Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling

TL;DR

The paper addresses the need to accelerate large-language-model pretraining by balancing learning-rate decay with batch-size growth through principled scheduling. It develops Seesaw, a batch-size scheduling algorithm that replaces LR halving with a 1/√2 LR factor while doubling the batch size, preserving loss dynamics and reducing serial steps. Theoretical results establish a non-asymptotic equivalence between LR decay and batch ramp-up for SGD on noisy linear regression and extend to normalized SGD as a proxy for Adam, with corollaries guiding practical scheduling. Empirically, Seesaw achieves ~36% wall-clock speedup at CBS across 150M/300M/600M models at Chinchilla scale, while matching cosine-annealing performance in FLOPs, and proves robust with AdamW, offering a practical, theory-grounded method for accelerating large-scale pretraining.

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.

Figures (8)

• Figure 1: Seesaw comparison with cosine decay in 150M (left), 300M (middle) and 600M (right) models trained at Chinchilla scale. Seesaw matches the loss dynamics of cosine annealing in FLOPs (top row), but achieves a significant speed up in terms of serial runtime (bottom row). Runs are swept over learning rates and plotted at the best learning rate for cosine annealing in terms of validation loss, at each batch size. The validation losses at the end of training are provided in Table \ref{['tab:final_losses']}. Note the axes: the top plots are on a logarithmic scale while the bottom are on a linear scale. For more experimental details, see Section \ref{['sec:exp_details']}.
• Figure 2: 150M models trained at batch size 256 (left) and 512 (right) with $\alpha$ and $\beta$ values following the line of equivalence $\alpha \sqrt{\beta} = 2$ described in Table \ref{['tab:between_points']}. Note that the target to match is the blue trace, and our theory (Lemma \ref{['lem:divergence_conditions']}) predicts that the red and purple traces should not match the baseline (blue trace) due to instabilities.
• Figure 3: 150M models trained past CBS (roughly 256), at batch sizes 1024, 2048, 4096 and 8192, for 3 schedulers: cosine decay (blue), constant learning rate with increasing batch size based on Seesaw (orange) and Seesaw (green). Note that none of the proposed schedules is able to match the cosine curve, with the discrepancy increasing as the batch size grows more.
• Figure 4: 150M experiments with weight decay across different batch sizes (128, 256, 512) for cosine annealing and Seesaw, for learning rate and weight decay values $(\eta, \lambda) = (0.003, 0.0001)$. Note that the losses overlap during training. We provide the final validation losses in Table \ref{['tab:weight_decay_final_losses']}.
• Figure 5: 150M models trained with 4 different schedules, at CBS (right) and just below (left). Blue trace keeps learning rate fixed and doubles batch size, orange trace keeps learning rate fixed and quadruples batch size, green trace halves learning rate at fixed batch size, and red trace is Seesaw. Note that the naive scheduling (blue) severely underperforms the baseline (green) and Seesaw (red).

Theorems & Definitions (13)

• Theorem : Informal version of Theorem \ref{['thm:sgd_bounded_risks']}
• Corollary : Informal version of Corollary \ref{['cor:normalized_sgd_bounded_risks']}
• Remark 1
• Lemma 1: Maximum Theoretical Speedup under Cosine Decay
• Theorem 1: SGD Equivalence
• Corollary 1: Normalized SGD Equivalence
• Lemma 2
• proof
• Lemma 3
• proof