Technical Summary: On the Nonlinearity of Learning Rate Scaling for LLM Training

Problem Statement

Training large language models (LLMs) efficiently requires careful hyperparameter tuning, yet directly sweeping learning rates at target scales is often prohibitively expensive. Consequently, practitioners rely on hyperparameter transfer, predicting optimal settings for large-scale models based on smaller-scale runs. A prevailing assumption in this domain is that the optimal learning rate ($\eta^*$) follows a log-linear (power-law) scaling relationship with respect to data scale ($D$) and model size ($N$):
$$\log \eta^*(D, N) = a \log D + b \log N + c$$
However, existing studies often fail to rigorously evaluate whether this log-linear assumption holds globally or only locally, and they frequently omit the computational cost required to fit these scaling laws. This paper investigates the validity of this scaling law and the efficiency of different extrapolation strategies (data-axis vs. model-axis) and parameterizations (learning rate vs. effective learning rate).

Methodology

The authors conducted a systematic empirical study using GPT-2-style architectures (ranging from 22M to 707M parameters) trained on the FineWeb-100B dataset with token budgets from 5B to 100B.

1. Experimental Design:

   • Models were trained using AdamW with a fixed batch size (0.52M tokens) and a warmup-steady-decay (WSD) schedule.
   • Learning rates were swept over a logarithmic grid ($\log_2 \eta \in [-16, -7]$).
   • Optimal Learning Rate ($\eta^*$): Estimated by fitting a cubic polynomial to the validation loss as a function of $\log \eta$.
   • Effective Learning Rate ($\eta_{eff}$): Defined as the step size in normalized weight space: $\eta_{eff}(t) = \|\hat{w}_{t+1} - \hat{w}_t\|_2$, where $\hat{w} = w/\|w\|_2$. The optimal effective learning rate ($\eta^*_{eff}$) was similarly estimated.

2. Evaluation Metrics:

   • Out-of-Distribution $R^2$ ($R^2_{OOD}$): Measures the predictive accuracy of scaling laws on unseen scales.
   • Extra Compute Ratio (ECR): Quantifies the additional computational cost incurred when using a predicted learning rate versus the true optimal one to achieve the same validation loss.
   • Compute-Constrained Extrapolation: The authors compared extrapolation along the data axis (varying $D$ with fixed large $N$) versus the model axis (varying $N$ with fixed large $D$) under identical compute budgets.

3. Theoretical Analysis:

   • The paper analyzes the implicit dynamics of AdamW, specifically the relationship between the raw learning rate $\eta$ and the effective learning rate $\eta_{eff}$, considering the transient behavior of weight norms.
   • Experiments were repeated using AdamH, an optimizer that explicitly controls $\eta_{eff}$, to isolate the effects of weight-norm dynamics.

Key Results

1. Nonlinearity of Optimal Learning Rate

The study finds that the optimal learning rate $\eta^*$ does not follow a global log-linear scaling law.

• Upward Curvature: As data scale $D$ and model size $N$ increase, $\eta^*$ exhibits systematic upward curvature. Log-linear extrapolation from smaller scales consistently underpredicts the optimal learning rate for larger scales.
• Linearity of Effective Learning Rate: In contrast, the optimal effective learning rate ($\eta^*_{eff}$) exhibits substantially more linear scaling behavior across the same ranges, with higher Pearson correlation coefficients.

2. Superiority of Data-Axis Extrapolation

When comparing extrapolation strategies under fixed compute budgets:

• Data-Axis vs. Model-Axis: Fitting scaling laws along the data-size axis ($D$) yields significantly more accurate predictions than along the model-size axis ($N$).
• Compute Efficiency: Using data-axis extrapolation with the effective learning rate incurs only ~2% additional compute relative to using the true optimal learning rate. This is less than 20% of the compute required for full hyperparameter sweeping at the target scale.

3. Mechanistic Explanation

The paper attributes the nonlinearity in $\eta^*$ to the transient dynamics of weight norms in AdamW:

• Two Regimes: The relationship between $\eta$ and $\eta_{eff}$ depends on whether the weight norm $\|w_t\|$ has reached equilibrium.
  • Equilibrium Regime: $\eta_{eff} \propto \sqrt{\eta}$ (slope 2 in log-log space).
  • Pre-equilibrium Regime: $\eta_{eff} \propto \eta$ (slope 1 in log-log space).
• Impact of Scale: As the optimal learning rate decreases for larger models/data, the training horizon often remains within the pre-equilibrium regime (where weight norms have not yet relaxed). This shifts the effective step size dynamics, causing the observed upward curvature in $\eta^*$.
• Validation via AdamH: When using AdamH to explicitly control $\eta_{eff}$, the scaling law becomes strictly log-linear, confirming that the nonlinearity in standard AdamW arises from the implicit, time-varying mapping between $\eta$ and $\eta_{eff}$.

Significance and Claims

The paper claims three primary contributions:

1. Empirical Correction: It demonstrates that the standard log-linear scaling law for learning rates is only a local property and fails at larger scales due to upward curvature, leading to inaccurate extrapolation.
2. Practical Transfer Strategies: It identifies two specific choices that substantially improve hyperparameter transfer:
   • Parameterizing the scaling law using the effective learning rate ($\eta_{eff}$) rather than the raw learning rate.
   • Extrapolating along the data axis ($D$) rather than the model size axis ($N$).
3. Theoretical Insight: It provides a mechanistic explanation for the observed nonlinearity, linking it to the transient behavior of weight norms in AdamW. The authors argue that $\eta_{eff}$ is the more fundamental quantity governing optimization dynamics in scale-invariant architectures.

The authors note limitations, including the focus on GPT-2 architectures and AdamW, the reliance on cubic fitting for $\eta^*$ estimation, and the analysis being limited to scales up to 100B tokens and 707M parameters. They do not claim these findings generalize automatically to all architectures or optimizers without further verification.