Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training

2025-05-19

AI Summary

This work addresses the scaling laws of hyperparameters—particularly learning rate, weight decay, and batch size—in large language model (LLM) pretraining with respect to model size (N), dataset size (D), and batch size (B). We propose a theoretical framework grounded in the time-scale invariance of AdamW: (B/(eta lambda D)). We establish that the optimal weight decay scales linearly with batch size ((lambda_{ ext{opt}} propto B)) and is precisely governed by a power law of the tokens-per-parameter ratio (D/N); furthermore, both the critical and optimal batch sizes scale as power laws of (D) alone, independent of (N). Through extensive controlled experiments and empirical validation of scaling exponents, we derive a Pareto-optimal joint selection strategy for (N) and (D) under dual objectives—training time and computational cost. The framework is validated across LLaMA and Pythia architectures, enabling high-accuracy *a priori* prediction of (lambda_{ ext{opt}}) and reducing hyperparameter tuning costs by orders of magnitude—e.g., millions of GPU-hours—thereby providing reusable, scalable configuration guidelines for thousand-GPU pretraining.

Technology Category

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

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Efficient LLM pre-training requires well-tuned hyperparameters (HPs), including learning rate {eta} and weight decay {lambda}. We study scaling laws for HPs: formulas for how to scale HPs as we scale model size N, dataset size D, and batch size B. Recent work suggests the AdamW timescale, B/({eta}{lambda}D), should remain constant across training settings, and we verify the implication that optimal {lambda} scales linearly with B, for a fixed N,D. However, as N,D scale, we show the optimal timescale obeys a precise power law in the tokens-per-parameter ratio, D/N. This law thus provides a method to accurately predict {lambda}opt in advance of large-scale training. We also study scaling laws for optimal batch size Bopt (the B enabling lowest loss at a given N,D) and critical batch size Bcrit (the B beyond which further data parallelism becomes ineffective). In contrast with prior work, we find both Bopt and Bcrit scale as power laws in D, independent of model size, N. Finally, we analyze how these findings inform the real-world selection of Pareto-optimal N and D under dual training time and compute objectives.

Problem

• Study scaling laws for hyperparameters in LLM pre-training
• Determine optimal weight decay and batch size scaling
• Analyze Pareto-optimal model and dataset size selection

Innovation

• Scaling laws for hyperparameters in LLM pre-training
• Optimal weight decay scales linearly with batch size
• Power law predicts optimal batch size independently of model size