Scaling Optimal LR Across Token Horizons

2024-09-30

AI Summary

Hyperparameter tuning of the learning rate (LR) in large language model (LLM) training is prohibitively expensive due to the scale of modern training regimes. Method: We conduct a systematic empirical study and scaling-law analysis to characterize how the optimal LR evolves with token horizon—the total number of training tokens—across diverse model sizes and architectures. Contribution/Results: We discover, for the first time, that the optimal LR follows a power-law decay with respect to token horizon. Based on this, we derive a zero-cost, transferable LR prediction rule that generalizes across training scales. Our rule enables accurate extrapolation of the optimal LR from short training runs to full-scale training, achieving <5% prediction error. Furthermore, applying it to LLaMA-1 reveals that its default LR is suboptimal, causing an approximately 0.8% increase in validation perplexity. This work provides both theoretical insight into LR scaling behavior and a practical tool for efficient, principled LLM training.

Technology Category

• Natural Language Processing: Learning & Optimization for NLP
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

State-of-the-art LLMs are powered by scaling -- scaling model size, dataset size and cluster size. It is economically infeasible to extensively tune hyperparameter for the largest runs. Instead, approximately optimal hyperparameters must be inferred or extit{transferred} from smaller experiments. Hyperparameter transfer across model sizes has been studied in Yang et al. However, hyperparameter transfer across dataset size -- or token horizon -- has not been studied yet. To remedy this we conduct a large scale empirical study on how optimal learning rate (LR) depends on token horizon in LLM training. We first demonstrate that the optimal LR changes significantly with token horizon -- longer training necessitates smaller LR. Secondly we demonstrate the the optimal LR follows a scaling law, and that the optimal LR for longer horizons can be accurately estimated from shorter horizons via such scaling laws. We also provide a rule-of-thumb for transferring LR across token horizons with zero overhead over current practices. Lastly we provide evidence that LLama-1 used too high LR, and estimate the performance hit from this. We thus argue that hyperparameter transfer across data size is an important and overlooked component of LLM training.

Problem

• Optimal learning rate varies with token horizon
• Scaling laws estimate LR for longer token horizons
• Hyperparameter transfer across data size is crucial

Innovation

• Scaling optimal learning rate
• Transfer hyperparameters across token horizons
• Empirical study on LLM training