📊 statistics

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

This paper introduces Seesaw, a principled framework that accelerates large language model pretraining by simultaneously scaling the batch size and adjusting the learning rate to preserve loss dynamics, thereby reducing wall-clock training time by approximately 36% compared to standard cosine decay schedules while matching performance at equal FLOPs.

Original authors: Alexandru Meterez, Depen Morwani, Jingfeng Wu, Costin-Andrei Oncescu, Cengiz Pehlevan, Sham Kakade

Published 2026-07-28

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Alexandru Meterez, Depen Morwani, Jingfeng Wu, Costin-Andrei Oncescu, Cengiz Pehlevan, Sham Kakade

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a giant, super-smart robot to understand the world by feeding it a mountain of books. This process is called "training," and it's the secret sauce behind the AI chatbots and tools we see today. To learn well, the robot needs two main things: a "learning rate," which is like how much it changes its mind after reading a single page, and a "batch size," which is how many pages it reads before stopping to think and adjust.

For a long time, scientists thought the best way to speed things up was to just feed the robot more pages at once (a bigger batch size) so it could process data faster. But there's a catch: if you feed it too many pages at once without changing how it thinks, it gets confused and stops learning efficiently. It's like trying to learn a language by reading a whole encyclopedia in one sitting; you might get the facts, but you won't understand the grammar. The big question researchers have been asking is: How do we balance reading more pages with thinking carefully enough to learn faster, without the robot getting lost?

This paper introduces a clever new strategy called Seesaw to solve that balancing act. The researchers discovered a mathematical rule that acts like a perfect seesaw: whenever you double the number of pages the robot reads at once (the batch size), you shouldn't just keep the learning rate the same or cut it in half. Instead, you should adjust the learning rate by a very specific amount—dividing it by the square root of two (about 1.41).

Think of it like this: if you double the size of your study group (batch size), you don't need to slow down your conversation speed (learning rate) as much as you might think. By using this specific "Seesaw" rhythm, the robot can process the same amount of information in fewer steps. The authors proved this mathematically for simple learning tasks and then tested it on massive language models with 150 million, 300 million, and 600 million parameters. They found that by using Seesaw, they could train these models about 36% faster in real-world time (wall-clock time) compared to the standard methods, while still achieving the exact same level of intelligence.

The paper also explicitly warns against being too greedy. If you try to increase the batch size too aggressively without adjusting the learning rate correctly, the robot's learning process becomes unstable and it stops improving. The authors showed that there is a "tipping point" where the math breaks down, and their Seesaw method stays safely on the right side of that line. In short, Seesaw isn't just a guess; it's a mathematically grounded recipe that lets AI models learn the same lessons in significantly less time, bringing us closer to building smarter AI without waiting months for the training to finish.

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

Problem Statement

Large Language Model (LLM) pretraining is increasingly constrained by wall-clock time, often extending to several months for state-of-the-art models. A primary strategy to reduce this duration is increasing the batch size, which leverages parallel computation to reduce the total number of sequential optimization steps. However, simply increasing the batch size indefinitely leads to diminishing returns beyond a "critical batch size" (CBS), where sample efficiency drops.

While recent large-scale training runs (e.g., LLaMA, Nemotron, OLMo) have employed "batch ramp" schedules—gradually increasing batch size over the course of training—these strategies are currently tuned heuristically. There is a lack of theoretical justification for how to optimally balance learning rate decay with batch size ramp-up, particularly for adaptive optimizers like Adam. The central question addressed is: What is the optimal batch size schedule to minimize serial runtime without sacrificing model performance?

Methodology

Theoretical Framework

The authors develop a principled framework based on noisy linear regression to establish an equivalence between learning rate decay and batch size ramp-up.

SGD Equivalence: The paper provides the first finite-sample proof (non-asymptotic) demonstrating that for Stochastic Gradient Descent (SGD) on noisy linear regression, a process where the batch size doubles every phase while the learning rate remains fixed is equivalent (up to a constant factor in excess risk) to a process where the batch size is fixed and the learning rate is halved every phase, provided both processes consume the same total number of data points.
Extension to Normalized SGD (NSGD): To bridge the gap to adaptive optimizers like Adam, the authors analyze Normalized SGD, a tractable analytical proxy for Adam. They introduce Assumption 2, which posits that in the variance-dominated regime observed in practice, the expected squared gradient norms are dominated by additive noise (scaling as $O(1/B)$ ). Under this assumption, NSGD effectively reduces to SGD with a rescaled learning rate.
The Equivalence Condition: The analysis derives a specific relationship required to maintain loss dynamics when adjusting both parameters. For a learning rate decay factor $\alpha$ and a batch size increase factor $\beta$ , the processes remain equivalent if the product $\alpha\sqrt{\beta}$ is held constant.

The Seesaw Algorithm

Based on the theoretical equivalence, the authors propose Seesaw, a scheduler that dynamically balances learning rate and batch size:

Mechanism: Whenever a standard scheduler (e.g., cosine decay) would reduce the learning rate by a factor of $\alpha$ , Seesaw instead reduces the learning rate by $\sqrt{\alpha}$ and increases the batch size by $\alpha$ .
Constraint: The authors identify a theoretical limit to prevent divergence. The most aggressive ramp-up scheme that maintains stability is defined by $\alpha = \sqrt{\beta}$ . Exceeding this (i.e., increasing batch size too aggressively relative to learning rate decay) leads to divergence.
Implementation: Seesaw is designed as a drop-in replacement for existing learning rate schedulers. In practice, it approximates cosine decay using step decay, triggering batch size increases at the specific token counts where the learning rate would otherwise be cut.

Key Contributions

Theoretical Equivalence: The paper establishes the first non-asymptotic equivalence between learning rate decay and batch size ramp-up for SGD and extends this to Normalized SGD under a variance-dominated regime.
The Seesaw Scheduler: Introduction of a principled algorithm that replaces standard learning rate decay with a coupled learning rate and batch size schedule, theoretically grounded to preserve loss dynamics.
Speedup Analysis: The authors derive a theoretical upper bound for the serial runtime reduction. By utilizing the most aggressive stable ramp-up ( $\alpha = \sqrt{\beta}$ ) against a cosine decay baseline, the maximum theoretical speedup is calculated as $(1 - 2/\pi) \approx 36.3\%$ .

Empirical Results

The authors evaluated Seesaw on 150M, 300M, and 600M parameter models trained at Chinchilla scale ( $D=20N$ ) using the OLMo codebase.

Performance Parity: Seesaw matches the validation loss dynamics of standard cosine annealing schedules across all model sizes when trained at the critical batch size.
Runtime Reduction: Seesaw achieves a significant reduction in serial runtime, reducing wall-clock time by approximately 36%, closely approaching the theoretical limit derived in the analysis.
Robustness: The method works effectively with AdamW and various weight decay settings.
Limits of Aggressiveness: Experiments varying the $\alpha$ and $\beta$ parameters confirm the theoretical constraint. Schedules that are more aggressive than $\alpha = \sqrt{\beta}$ (e.g., doubling the batch size while keeping the learning rate constant or decaying it too slowly) result in divergence or suboptimal performance.
Critical Batch Size Limit: The authors observe that the Seesaw strategy is effective up to the critical batch size. Beyond this point, where the noise term no longer dominates the gradient norm, the equivalence breaks down, and Seesaw performs worse than standard cosine decay.

Significance and Claims

The paper claims to provide a rigorous theoretical foundation for the "batch ramp" practice, moving it from heuristic tuning to a principled optimization strategy. By demonstrating that learning rate decay and batch size ramp-up are interchangeable under specific conditions, Seesaw offers a practical method to accelerate LLM pretraining without compromising model quality.

The authors emphasize that Seesaw serves as a drop-in replacement for existing schedulers in adaptive optimizers. The primary significance lies in the ability to reduce the wall-clock time of training large models by roughly one-third, addressing the hardware bottleneck where training requirements are outpacing hardware improvements. The work also clarifies the boundaries of this strategy, noting that it relies on the variance-dominated regime and may not apply once batch sizes exceed the critical threshold where mean gradient terms dominate.