TLDR: Seesaw is a new method for training large language models faster. It intelligently adjusts both the learning rate and batch size during training, specifically by doubling the batch size and slightly reducing the learning rate when a standard scheduler would halve the learning rate. This approach is theoretically grounded and empirically shown to reduce LLM pretraining wall-clock time by about 36% without sacrificing performance, especially for adaptive optimizers like Adam.
Training large language models (LLMs) is a monumental task, often taking months to complete. This extensive duration is a significant bottleneck in the rapid advancement of AI. Researchers are constantly seeking ways to accelerate this process without compromising the models’ performance. One promising strategy involves gradually increasing the batch size during training, a technique often referred to as a “batch ramp.” However, the optimal way to implement these batch ramps, especially with adaptive optimizers like Adam, has largely been based on trial-and-error heuristics rather than solid theoretical principles.
A new research paper introduces a principled framework for batch-size scheduling called Seesaw. This innovative approach aims to accelerate LLM pretraining by intelligently balancing the learning rate and batch size throughout the training process. The core idea behind Seesaw is elegant: whenever a standard learning rate scheduler would typically halve the learning rate, Seesaw instead multiplies the learning rate by 1/√2 and simultaneously doubles the batch size. This specific adjustment is designed to maintain the crucial loss dynamics of the training process while significantly reducing the number of sequential steps, thereby cutting down wall-clock time.
The Science Behind Seesaw
The theoretical underpinnings of Seesaw are robust. The paper provides what is believed to be the first finite-sample proof demonstrating the equivalence between learning-rate decay and batch-size ramp-up for Stochastic Gradient Descent (SGD) in the context of noisy linear regression. This equivalence is further extended to normalized SGD, which serves as a tractable proxy for adaptive optimizers like Adam, particularly under a variance-dominated regime commonly observed in practice.
Informally, for SGD, the theory suggests that doubling the batch size can be equivalent to halving the learning rate. For normalized SGD, the relationship is more nuanced: a learning rate cut by a factor of ‘alpha’ and a batch size increase by a factor of ‘beta’ are considered equivalent as long as the product of ‘alpha’ and the square root of ‘beta’ remains constant. Seesaw leverages this relationship, specifically choosing parameters that allow for aggressive batch size increases while maintaining stable training dynamics.
Empirical Validation and Speedup
The practical implications of Seesaw are impressive. Empirical evaluations on large language models with 150 million, 300 million, and 600 million parameters, trained at Chinchilla scale, show that Seesaw effectively matches the performance of traditional cosine decay schedules in terms of computational effort (FLOPs). Crucially, it achieves this while reducing the wall-clock training time by approximately 36%. This reduction closely approaches the theoretical limit predicted by the paper’s analysis.
Furthermore, Seesaw has been shown to work effectively even when used with AdamW, a popular adaptive optimizer, and tuned weight decay, making it a practical and versatile solution for current LLM pretraining workflows. The method acts as a direct replacement for existing cosine schedulers, offering a significant speedup without requiring complex re-tuning of other hyperparameters.
Also Read:
- Understanding LLM Performance: A New Framework for Context and Compute Scaling
- Unlocking Predictable Scaling for Reinforcement Learning in Large Language Models
Understanding the Limits
While Seesaw offers substantial benefits, the research also explores its boundaries. A key assumption for Seesaw’s effectiveness is that the expected gradient norms are dominated by additive noise. This assumption holds true for a wide range of training scenarios. However, at extremely large batch sizes, beyond what is known as the critical batch size (CBS), this assumption can break down. In such scenarios, the noise variance becomes very small, and simply increasing the batch size further does not yield the same benefits. The paper demonstrates that past a certain batch size, learning rate decay becomes indispensable for achieving lower losses, and batch size ramp-up alone is insufficient to match the performance of learning rate decay.
In conclusion, Seesaw represents a significant step forward in optimizing LLM training. By providing a theoretically grounded and empirically validated method for dynamically scheduling learning rates and batch sizes, it offers a practical way to accelerate the pretraining of large language models, making advanced AI development more efficient. For more details, you can read the full research paper here.
