Reproduce “Scaling Optimal LR Across Token Horizons” on OpenResearch

About this paper

Johan Bjorck, Alon Benhaim, Vishrav Chaudhary, Furu Wei, Xia Song

This research from Nvidia, Microsoft, and Meta empirically investigates how the optimal learning rate for Large Language Model (LLM) training scales with the total number of training tokens, demonstrating that it decreases following a power-law relationship. The study provides a methodology to predict optimal learning rates for very long training runs from shorter experiments, illustrating its utility by showing Llama-1's original training likely used a suboptimal learning rate, incurring an estimated 0.027 validation loss penalty.

Key insights

• The optimal learning rate for LLM training consistently decreases as the token horizon lengthens, a phenomenon that `muP` parametrization, designed for model size transfer, does not inherently address.
• This decrease in optimal learning rate adheres to a robust power-law scaling relationship with the token horizon, exhibiting high R^2 values (0.96 to 0.99) across diverse model sizes.
• The exponent of the optimal learning rate's power-law scaling with token horizon (β) is approximately constant across different model sizes and largely independent of batch size, suggesting generalizability and a factorizable dependence.

Summary by alphaXiv.