Summary of "On the Nonlinearity of Learning Rate Scaling for LLM Training"

Summary (Overview)

• Key finding: The optimal learning rate $\eta^{*}$ for LLM training exhibits upward curvature at larger scales, deviating from the commonly assumed log-linear (power-law) scaling law, leading to underpredictions when extrapolating from smaller runs.
• Effective learning rate as a better parameterization: The effective learning rate $\eta_{\mathrm{eff}}$ (step size in normalized weight space) demonstrates substantially more linear scaling behavior than the raw learning rate, making it a more reliable quantity for hyperparameter transfer.
• Data-axis beats model-axis extrapolation: Fitting scaling laws along the data size axis $D$ yields much more accurate predictions than along the model size axis $N$, requiring only ~2% additional training compute relative to using the true optimal learning rate.
• Mechanistic explanation: The nonlinearity is explained via weight-norm dynamics—small learning rates keep training in a pre-equilibrium regime where the $\eta \leftrightarrow \eta_{\mathrm{eff}}$ mapping has slope 1, while larger learning rates reach equilibrium with slope 2, causing the curvature.
• Validation with AdamH: Explicitly controlling $\eta_{\mathrm{eff}}$ via the AdamH optimizer recovers clean log-linear scaling, confirming the theoretical explanation.

Introduction and Theoretical Foundation

Background and Motivation

Training large language models requires carefully tuned hyperparameters, but sweeping these directly at target scale is prohibitively expensive. Hyperparameter transfer—predicting effective hyperparameters at large scale from inexpensive smaller-scale runs—has become a central tool in modern model development.

Two main lines of work exist:

1. Maximal Update Parameterization ($\mu$P): Theoretically motivated framework ensuring optimal learning rate remains stable across width scaling, though it focuses on model scaling and doesn't cover data-scale transfer.
2. Power-law scaling laws: Directly modeling the optimal learning rate $\eta^{*}$ as a function of data scale $D$ and model size $N$.

Theoretical Foundation

The standard assumption is a log-linear (power-law) scaling law:

$$\log \eta^{*}(D, N) = a \log D + b \log N + c \tag{2}$$

which implies $\eta^{*}(D, N) \propto D^{a}N^{b}$.

The effective learning rate is defined as the step size in normalized weight space:

$$\eta_{\mathrm{eff}}(t) = \|\hat{w}_{t+1} - \hat{w}_{t}\|_{2} \tag{4}$$

where $\hat{w} = w/\|w\|_{2}$ denotes the normalized weight direction. This is motivated by scale-invariance of modern architectures: for weights $w$ with normalization layers, $\mathcal{L}(w) = \mathcal{L}(\alpha w)$ for all $\alpha > 0$ (Equation 3).

The optimal learning rate is defined as:

$$\eta^{*}(D, N) := \arg\min_{\eta} \mathcal{L}\left(w_{T}(\eta; D, N)\right) \tag{1}$$

and the optimal effective learning rate as:

$$\eta_{\mathrm{eff}}^{*}(D, N) := \arg\min_{\overline{\eta}_{\mathrm{eff}}} \mathcal{L}\left(\hat{w}_{T}(\overline{\eta}_{\mathrm{eff}}; D, N)\right) \tag{6}$$

Methodology

Experimental Setup

• Models: GPT-2 style models from 22M to 707M parameters
• Data: FineWeb-100B with data budgets from 5B to 100B tokens in 2.5B-token increments
• Optimizer: AdamW with fixed hyperparameters ($\beta_1, \beta_2, \epsilon$, weight decay $\lambda = 0.1$), batch size 0.52M tokens
• Learning rate sweep: Logarithmic grid with $\log_2 \eta \in [-16, -7]$ with unit increments
• Schedule: Warmup-steady-decay (WSD) schedule

Scaling Laws Evaluated

Four log-linear scaling laws were evaluated:

For optimal learning rate:

• Fixed model size $N$: $\log \eta^{*}(D, N) = a_N \log D + b_N$ (Eq. 7)
• Fixed data size $D$: $\log \eta^{*}(D, N) = c_D \log N + d_D$ (Eq. 8)

For optimal effective learning rate:

• Fixed model size $N$: $\log \eta_{\mathrm{eff}}^{*}(D, N) = \tilde{a}_N \log D + \tilde{b}_N$ (Eq. 9)
• Fixed data size $D$: $\log \eta_{\mathrm{eff}}^{*}(D, N) = \tilde{c}_D \log N + \tilde{d}_D$ (Eq. 10)

Evaluation Metrics

Out-of-distribution coefficient of determination ($R^2_{OOD}$):

$$R_{\mathrm{OOD}}^{2}(\beta) = 1 - \frac{\sum_{(D,N) \in \mathcal{S}_{\mathrm{test}}} \left(\log \eta^{*}(D,N) - \log \hat{\eta}_{\beta}^{*}(D,N)\right)^{2}}{\sum_{(D,N) \in \mathcal{S}_{\mathrm{test}}} \left(\log \eta^{*}(D,N) - \overline{\log \eta^{*}}\right)^{2}} \tag{11}$$

Extra Compute Ratio (ECR):

$$\mathrm{ECR}(\beta) = \frac{\sum_{(D,N) \in \mathcal{S}_{\mathrm{test}}} \Delta C(D,N)}{\sum_{(D,N) \in \mathcal{S}_{\mathrm{test}}} C(D,N)} \tag{12}$$

where $C = 6ND$ (FLOPs for a dense Transformer) and $\Delta C$ is the additional compute needed for the predicted learning rate to match the optimal validation loss.

Compute-Constrained Training Sets

Two extrapolation strategies were tested (both with total compute budget $\beta C(D_{\mathrm{test}}, N_{\mathrm{test}})$):

• D-axis extrapolation: $S_{\mathrm{train}}^{(D)} = \{(\alpha_k D_{\mathrm{test}}, N_{\mathrm{test}})\}_{k=1}^K$
• N-axis extrapolation: $S_{\mathrm{train}}^{(N)} = \{(D_{\mathrm{test}}, \alpha_k N_{\mathrm{test}})\}_{k=1}^K$

with constraint $\sum_k \alpha_k = \beta$.

Empirical Validation / Results

Deviation from Log-Linear Scaling

Figure 2 shows that $\eta^{*}$ vs. $D$ and $N$ in log-log space exhibits systematic upward curvature at larger scales, indicating log-linearity is only a local property.

Figure 3 shows that $\eta_{\mathrm{eff}}^{*}$ exhibits improved linearity with higher Pearson correlation coefficients, particularly at larger scales.

Quantitative Extrapolation Performance

Figure 4a compares $R^2_{OOD}$ across compute budgets:

• Scaling along D consistently outperforms scaling along N
• $\eta_{\mathrm{eff}}^{*}$ consistently achieves higher $R^2_{OOD}$ than $\eta^{*}$ for both D-dim and N-dim fits

Figure 4b reveals a linear relationship between extrapolation error and compute cost:

$$\mathrm{ECR} \propto 1 - R_{\mathrm{OOD}}^{2} \tag{13}$$

Table 1: Concrete Case Study (Test: $D = 100B$, $N = 707M$)

Bold indicates the lowest ECR within each $\beta$ block. The D-dim + $\eta_{\text{eff}}$ setup consistently achieves the lowest extra compute ratio.

Mechanistic Explanation: Two Learning-Rate Phases

The evolution of $\eta_{\mathrm{eff}}$ under AdamW with constant $\eta$ is:

$$\eta_{\mathrm{eff}}(t) = \sqrt{2\eta\lambda \frac{1-\beta_1}{1+\beta_1} \frac{1}{1 + \left(\frac{W_0^2}{W_{+\infty}^2} - 1\right) e^{-2\eta\lambda t}}} \tag{15}$$

where $W_0 = \|w_0\|_2$ is the initial weight norm and $W_{+\infty} = U\sqrt{\frac{\eta}{2\lambda}\frac{1+\beta_1}{1-\beta_1}}$ is the equilibrium weight norm.

This yields two asymptotic regimes:

1. Equilibrium regime ($t \gg (2\eta\lambda)^{-1}$):

$$\eta_{\mathrm{eff}} \propto \sqrt{\eta}, \quad \log \eta = 2\log \eta_{\mathrm{eff}} + \text{const} \tag{16}$$

2. Pre-equilibrium regime ($t \ll (2\eta\lambda)^{-1}$):

$$\eta_{\mathrm{eff}} \propto \eta, \quad \log \eta = \log \eta_{\mathrm{eff}} + \text{const} \tag{17}$$

Key insight: Smaller learning rates have slower weight-norm convergence (decay time $(2\eta\lambda)^{-1}$), keeping training in the pre-equilibrium regime with slope-1 mapping. As training horizon grows, the optimal learning rate decreases, pushing training from equilibrium into pre-equilibrium regime—explaining the upward curvature.

Validation with AdamH

AdamH, which explicitly controls $\eta_{\mathrm{eff}}$ by normalizing both weight and update to fixed norms, recovers:

• Clean log-linear relationship: $\log \eta_{\mathrm{eff}}^{*} \propto \log D$ (Pearson $r = -0.9994$)
• Much smaller ECR than AdamW at small compute-budget ratios $\beta$

Theoretical and Practical Implications

Theoretical Implications

1. Log-linear scaling is only locally valid: The assumption that $\eta^{*} \propto D^a N^b$ holds globally is incorrect; systematic curvature emerges at larger scales.

2. Effective learning rate is the fundamental quantity: Due to scale-invariance of modern architectures, $\eta_{\mathrm{eff}}$—not $\eta$—governs optimization dynamics, and it is the quantity for which clean scaling laws hold.

3. Weight-norm dynamics matter: The transient behavior of weight norms under AdamW creates a time-varying mapping between $\eta$ and $\eta_{\mathrm{eff}}$, which is the root cause of nonlinearity.

Practical Implications

1. Recommended transfer strategy: Fit scaling laws along the data axis using the effective learning rate—this incurs only ~2% additional training compute relative to the true optimal learning rate.

2. Compute efficiency: D-axis extrapolation requires less than 20% of the compute of hyperparameter sweeping at the target scale.

3. Cost of extrapolation error: The linear relationship $\mathrm{ECR} \propto 1 - R^2_{OOD}$ means even modest extrapolation degradation directly translates to proportional compute waste.

Conclusion

Main Takeaways

1. Nonlinearity discovered: The optimal learning rate exhibits systematic upward curvature at larger scales, breaking the commonly assumed log-linear scaling law.
2. Two improvements identified: Parameterizing by effective learning rate and extrapolating along the data axis substantially improve transfer accuracy.
3. Mechanistic explanation provided: The curvature arises from weight-norm transient dynamics under AdamW, where the $\eta \leftrightarrow \eta_{\mathrm{eff}}$ mapping shifts between slope-1 (pre-equilibrium) and slope-2 (equilibrium) regimes.
4. AdamH validation: Explicitly controlling $\eta_{\mathrm{eff}}$ recovers clean linear scaling, confirming the theoretical framework.

Limitations

• Model/optimizer coverage: Only GPT-2-style architectures with AdamW on FineWeb were tested; generalization to other settings remains unverified.
• Estimation bias: Cubic polynomial fitting for $\eta^{*}$ may introduce bias near sweep boundaries.
• Scale limits: Analysis limited to 100B tokens and 707M parameters; more extreme regimes may exhibit additional phenomena.

Future Directions

• Extending the analysis to other architectures, optimizers, and more extreme scales
• Developing practical tools that directly control effective learning rate during training
• Further theoretical analysis of the transient dynamics and their impact on scaling behavior