On the Nonlinearity of Learning Rate Scaling for LLM Training

2026-06-27

AI Summary

Existing learning rate scaling laws often fail to accurately extrapolate to large-scale training due to their neglect of nonlinear effects. This work systematically investigates the training dynamics of GPT-2 models (ranging from 22M to 707M parameters) trained on datasets spanning 5B to 100B tokens, revealing that the optimal learning rate exhibits a concave nonlinear dependence on both model size and data volume. To address this, the authors propose a novel scaling strategy based on effective learning rate and dataset size, supported by a theoretical analysis grounded in the convergence dynamics of weight norms. The approach significantly improves extrapolation accuracy across diverse scale configurations, thereby reducing hyperparameter tuning costs for large models. Its robustness is further validated through experiments employing the AdamH optimizer and weight-space normalization techniques.

Abstract

Learning-rate transfer can reduce the cost of training large language models: instead of sweeping learning rates at target scale, practitioners extrapolate from smaller runs. Existing approaches often assume that the optimal learning rate follows a log-linear scaling law in data scale and model size. We carefully examine and evaluate this scaling law. In our empirical study of GPT-2--style models from 22M to 707M parameters trained on 5B to 100B tokens, the optimal learning rate develops upward curvature at larger scales, leading to inaccurate extrapolation. We find that this curvature largely disappears when learning rates are replaced by effective learning rate (the step size in normalized weight space), and when data $D$ extrapolation is used instead of model size $N$ extrapolation. Next, we explain nonlinearity in scaling: weight-norm converges to equilibrium slower when optimal learning is small, requiring a larger step size to reduce the transient phase. Experiments with AdamH, which directly controls the effective learning rate, further support this explanation.

Problem

• learning rate scaling
• large language models
• nonlinearity
• extrapolation
• effective learning rate

Innovation

• learning rate scaling
• effective learning rate
• nonlinear scaling
• large language models
• AdamH