Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling (2510.14717v1)

Abstract: Increasing the batch size during training -- a ''batch ramp'' -- is a promising strategy to accelerate LLM 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.

• The paper introduces the Seesaw algorithm that jointly schedules learning rate decay and batch size ramp-up to reduce wall-clock training time by approximately 36%.
• The paper provides a finite-sample proof demonstrating that halving the learning rate is equivalent to doubling the batch size under variance-dominated regimes, supporting the proposed method.
• The paper empirically validates Seesaw's robustness across different model scales and optimizer settings, showing comparable loss dynamics to standard schedulers with enhanced runtime efficiency.

Seesaw: A Principled Approach to Accelerating LLM Training via Joint Learning Rate and Batch Size Scheduling

Motivation and Problem Statement

The exponential growth in pretraining compute for LLMs has led to significant increases in wall-clock training time, often stretching to months for state-of-the-art models. While increasing batch size is a well-established method for reducing serial runtime, its effectiveness is bounded by the critical batch size (CBS), beyond which sample efficiency degrades. Recent large-scale LLM training runs have adopted heuristic batch size ramping schedules, but these lack theoretical justification and may not be optimal. This paper addresses the question: What is the optimal batch size schedule for minimizing serial runtime without sacrificing model performance?

Theoretical Framework: Equivalence of Learning Rate Decay and Batch Size Ramp

The central theoretical contribution is a finite-sample proof of equivalence between learning rate decay and batch size ramp-up for SGD in noisy linear regression. The analysis is extended to normalized SGD (NSGD), a tractable proxy for Adam, under a variance-dominated regime. The key result is that, under certain assumptions, halving the learning rate is equivalent to doubling the batch size, up to constant factors in excess risk. For NSGD, the equivalence generalizes: any learning rate cut by $\alpha$ and batch size increase by $\beta$ are equivalent as long as $\alpha \sqrt{\beta}$ is held constant.

This leads to the Seesaw algorithm, which, at each scheduled learning rate decay, multiplies the learning rate by $1/\sqrt{2}$ and doubles the batch size, preserving loss dynamics while reducing the number of serial steps.

Algorithmic Implementation: Seesaw Scheduler

The Seesaw algorithm is designed as a drop-in replacement for standard learning rate schedulers (e.g., cosine annealing) in adaptive optimizers. The implementation is straightforward:

def seesaw_scheduler(eta_0, B_0, alpha, S, T): eta = eta_0 B = B_0 for t in range(1, T+1): if t in S: eta /= np.sqrt(alpha) B *= alpha # optimizer step with current eta, B

Where $S$ is the array of steps at which the scheduler would cut the learning rate by $\alpha$.

Empirical Results: Runtime Acceleration and Loss Dynamics

Empirical evaluation is conducted on 150M, 300M, and 600M parameter models at Chinchilla scale ($D = 20N$), using the OLMo codebase and AdamW optimizer. Seesaw matches the loss dynamics of cosine annealing in terms of FLOPs, but achieves a serial runtime reduction of approximately 36%, closely approaching the theoretical limit derived in the analysis.

Figure 1: Seesaw matches the loss dynamics of cosine annealing in FLOPs (top row), but achieves significant speedup in serial runtime (bottom row) across 150M, 300M, and 600M models.

Further experiments validate the theoretical constraint that the most aggressive ramp-up scheme is $\alpha = \sqrt{\beta}$; more aggressive schedules lead to instability and divergence.

Figure 2: Loss dynamics for 150M models at batch sizes 256 and 512, with $\alpha$ and $\beta$ values along the equivalence line $\alpha \sqrt{\beta} = 2$; aggressive schedules (red, purple) underperform the baseline (blue).

When batch size exceeds CBS, Seesaw and other ramping schedules fail to match the performance of cosine decay, with the discrepancy increasing as batch size grows.

Figure 3: At batch sizes past CBS, none of the proposed schedules match the cosine curve, and the gap widens with increasing batch size.

Robustness to Weight Decay and Auxiliary Losses

Seesaw is robust to the inclusion of weight decay in AdamW, with overlapping loss curves across batch sizes and schedulers.

Figure 4: Losses for cosine annealing and Seesaw overlap during training with weight decay across batch sizes.

Ablations on Z-Loss show no difference in final validation loss at 150M scale, though some instabilities are observed at larger scales with Seesaw.

Figure 5: Final validation losses are equal whether Z-Loss is enabled or not for cosine decay at Chinchilla scale.

Figure 6: 600M models trained with Seesaw decay in Chinchilla scale, with Z-Loss.

Comparison to Alternative Schedulers

Naive scheduling strategies, such as keeping the learning rate fixed while increasing batch size, severely underperform compared to Seesaw and standard learning rate decay.

Figure 7: Naive scheduling (blue) underperforms the baseline (green) and Seesaw (red) at and just below CBS.

Theoretical Implications and Limitations

The equivalence between learning rate decay and batch size ramp is rigorously established for SGD and NSGD under variance-dominated regimes. However, the analysis breaks down when the additive noise assumption fails, i.e., at very large batch sizes. In such regimes, learning rate decay becomes essential for escaping stable cycles and achieving lower loss, as demonstrated by a toy NGD example.

Practical Implications and Future Directions

Seesaw provides a principled, theoretically justified method for accelerating LLM pretraining, reducing wall-clock time without compromising performance. The algorithm is simple to implement and can be integrated into existing training pipelines with minimal modification. The approach is robust to optimizer choice (AdamW, NSGD) and regularization (weight decay).

Future work should investigate the behavior of Seesaw at scales where the variance-dominated assumption fails, explore its interaction with auxiliary losses (e.g., Z-Loss) at larger model sizes, and extend the theoretical analysis to more complex loss landscapes and non-quadratic regimes.

Conclusion

This paper establishes a formal equivalence between learning rate decay and batch size ramping in SGD and NSGD, leading to the Seesaw algorithm for joint scheduling. Empirical results demonstrate that Seesaw matches the performance of standard schedulers while reducing serial runtime by approximately 36%. The findings provide a rigorous foundation for batch ramping practices in LLM training and offer a practical solution for accelerating pretraining at scale. Theoretical and empirical limitations at extreme batch sizes highlight avenues for further research in optimization dynamics and scheduler design.