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

TL;DR

Seesaw is a principled drop-in replacement for cosine learning-rate decay that interleaves LR cuts with batch-size doublings: whenever a standard scheduler would halve the LR, Seesaw multiplies the LR by 1/√2 and doubles the batch size instead. The authors prove (the first finite-sample result of this kind) that learning-rate decay and batch-size ramp-up are risk-equivalent for SGD on noisy linear regression, and extend the equivalence to normalized SGD as a tractable proxy for Adam under a variance-dominated regime. Empirically at 150M / 300M / 600M-parameter Chinchilla-scale pretraining on C4, Seesaw matches cosine decay at equal FLOPs while reducing serial (wall-clock) steps by ≈36%, approaching the theoretical maximum implied by the analysis. The recipe is preserved under AdamW with tuned weight decay, making it a practical accelerator for at-or-below-critical-batch-size pretraining runs.

Key claims

For SGD on noisy linear regression, a base process that doubles batch size after each phase at fixed LR has excess risk within a constant factor of an alternative process that halves LR at fixed batch size, processing the same data per phase [Theorem 1, §5.1].
The equivalence extends to normalized SGD (a tractable proxy for Adam) under a “variance-dominated” assumption that the additive-noise term dominates E[‖g_t‖²]; the equivalence becomes: cutting LR by α is equivalent to scaling batch size by 1/α² so long as B · η² is held constant [Corollary 1, §5.1; Eq. 7].
Seesaw is the practical algorithm derived from Corollary 1: at every point where a step-decay approximation of cosine would halve the LR, multiply LR by 1/√2 and double batch size [Algorithm 1, §3.2].
The most aggressive non-divergent ramp under the B·η² = const family is η_k = (1/√2)^k, B_k = 2^k; ramping faster (e.g. η_k = (1/2)^k, B_k = 4^k) empirically underperforms cosine [Lemma 4, §4.1, Fig. 2, Table 2].
At critical batch size on 150M / 300M / 600M models pretrained at Chinchilla scale (D ≈ 20·N tokens) on C4 with the OLMo codebase, Seesaw matches cosine-decay final loss across initial batch sizes 128 / 256 / 512 / 1024 within ≤0.012 nats [Table 1].
Theoretical maximum serial-runtime reduction under Seesaw vs cosine in the continuous limit is 1 − π/(2(π−1)) ≈ 26.6%; empirically Seesaw achieves ≈36% wall-clock reduction at Chinchilla scale [Lemma 1, §3.2, Fig. 1].
The “variance-dominated” Assumption 2 fails at sufficiently large batch — once batch size shrinks the noise term below the deterministic-gradient term, neither the NSGD-equivalent (B·η² const) nor the SGD-equivalent (B·η const) ramp matches cosine, and LR decay becomes irreducibly necessary [§4.2, Fig. 3].
Seesaw is robust to AdamW with weight decay tuned over {0, 0.01, 0.1} and remains equivalent to cosine at the same FLOPs [Appendix C, Fig. 4].

Method

Seesaw operationalizes a risk-equivalence between two phase-indexed schedules under (normalized) SGD. The theoretical setup considers two procedures that each consume the same number of samples per phase, where in phase k one schedule uses (η, B·2^k) (fixed LR, ramping batch) and the other uses (η/2^k, B) (decaying LR, fixed batch). The paper proves that the population risks at the end of each phase differ by at most a constant factor (Theorem 1 for SGD; Corollary 1 for NSGD under the variance-dominated assumption that E[‖g_t‖²] ≈ σ²/B). The Adam analysis reduces to NSGD by treating √(v_t) as the population-mean preconditioner — a standard tractable proxy used in earlier µP and SDE work.

The practical algorithm approximates cosine decay with a step-decay scheduler that cuts LR by 1/2 at the same token counts where cosine would (after a 10% warmup), then replaces each LR ← LR/2 event with the joint update LR ← LR/√2, B ← 2·B. The maximum aggressiveness is bounded by Lemma 4, which derives a divergence condition on the family of equivalent schedules (η_k = η₀·α^k, B_k = B₀·α^{−2k}) and locates the boundary at α = 1/√2. Above this boundary the variance bound on the iterates blows up.

Results

Headline serial-runtime reduction: ≈36% wall-clock speedup vs cosine at equal FLOPs and matching final loss, at 150M / 300M / 600M parameters and Chinchilla token budgets, with the theoretical-maximum 26.6% serial-step reduction in the continuous limit recovered (and slightly exceeded due to the step-decay vs continuous-cosine gap) [§3.2, §4, Fig. 1].
Final-loss equality with cosine across all model sizes and initial batch sizes tested, with deltas under ±0.012 nats and Seesaw winning slightly more often than losing — i.e. no quality regression [Table 1].
The most-aggressive ramp α = 1/2 (LR halves, batch quadruples) diverges or underperforms cosine sharply; α = 1/√2 is the cleanest predictor and the limit of the equivalence [§4.1, Fig. 2].
Beyond a model-size-dependent batch threshold (~512 for 150M, ~1024 for 600M), Assumption 2 breaks and neither Seesaw nor the SGD-rule (B·η = const) matches cosine — the paper’s NGD-on-quadratic toy argument shows LR decay is then irreducibly required [§4.2, Fig. 3].

Why it’s interesting

Seesaw is the first theoretically-grounded answer to a question the wiki’s hyperparameter-scaling cluster has had open for weeks: is “batch size as a free hyperparameter” robust across optimizers, or specific to the AdamW + CompleteP + SDE recipe of The Design Space of Tri-Modal Masked Diffusion Models? The paper’s B·η² = const equivalence is a strictly different rule than the AdamW-SDE result (which decouples physical from logical batch size up to a critical threshold without any LR cut) — and it provides a closed-form non-asymptotic proof rather than an SDE-limit argument. The two results are complementary: SDE-reparametrization tells you what physical-vs-logical batch sizes are interchangeable at fixed schedule; Seesaw tells you which (LR, B) trajectories through schedule space are interchangeable. Both undercut the existing practice in How to Set the Learning Rate for Large-Scale Pre-training? of treating batch size as a separately-tuned axis of the Fitting paradigm. The analysis also rules out the more aggressive batch-warmup rule of Merrill et al. (2025) (proposed B²·η = const or similar) — Lemma 4’s divergence boundary at α = 1/√2 says that family will eventually diverge, which is the cleanest theoretical critique of an existing heuristic in the cluster.