A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules

2025-03-17

AI Summary

Existing studies lack a quantitative model characterizing how learning rate schedules affect loss evolution during large language model pretraining. Method: We propose the Multi-Power-Law Loss Evolution Law—a unified analytical framework that models loss curves under canonical schedules (e.g., constant, cosine, stepwise) via decoupled components: the cumulative effect of learning rate and the additional convergence gain induced by decay. This enables cross-schedule generalization and accurate prediction of unseen schedule shapes or durations from minimal fitting data. Contribution/Results: Leveraging this law, we automatically discover a novel Warmup-Stable-Decay schedule structure, which consistently achieves significantly lower final loss than standard cosine annealing across diverse model scales and architectures. The model provides an interpretable, optimization-friendly theoretical tool for efficient pretraining—requiring only sparse schedule-loss observations for high-fidelity loss curve prediction.

Technology Category

• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Learning to Search
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Training large models is both resource-intensive and time-consuming, making it crucial to understand the quantitative relationship between model performance and hyperparameters. In this paper, we present an empirical law that describes how the pretraining loss of large language models evolves under different learning rate schedules, such as constant, cosine, and step decay schedules. Our proposed law takes a multi-power form, combining a power law based on the sum of learning rates and additional power laws to account for a loss reduction effect induced by learning rate decay. We extensively validate this law on various model sizes and architectures, and demonstrate that after fitting on a few learning rate schedules, the law accurately predicts the loss curves for unseen schedules of different shapes and horizons. Moreover, by minimizing the predicted final pretraining loss across learning rate schedules, we are able to find a schedule that outperforms the widely used cosine learning rate schedule. Interestingly, this automatically discovered schedule bears some resemblance to the recently proposed Warmup-Stable-Decay (WSD) schedule (Hu et al, 2024) but achieves a slightly lower final loss. We believe these results could offer valuable insights for understanding the dynamics of pretraining and designing learning rate schedules to improve efficiency.

Problem

• Predicts pretraining loss evolution across learning rate schedules.
• Validates multi-power law on various model sizes and architectures.
• Optimizes learning rate schedules to minimize final pretraining loss.

Innovation

• Multi-power law predicts loss curves
• Validated across model sizes and architectures
• Automatically discovers optimal learning rate schedules