# On the Nonlinearity of Learning Rate Scaling for LLM Training

> The optimal learning rate for LLM training exhibits upward curvature at scale, breaking log-linear scaling, but effective learning rate with data-axis extrapolation transfers accurately with only ~2% extra compute.

- **Source:** [arXiv](https://arxiv.org/abs/2606.29158)
- **Published:** 2026-08-22
- **Permalink:** https://picx.dev/p/iETSYd
- **Whiteboard:** https://picx.dev/p/iETSYd/image

## Summary

# 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$)

| Budget Ratio (%) | Setup | Log LR Prediction Error | Val Loss Gap ($10^{-4}$) | ECR (%) |
|:---:|:---:|:---:|:---:|:---:|
| 13.00 | D-dim, $\eta_{\text{eff}}$ | +0.1235 | 4.77 | **1.21** |
| 13.00 | D-dim, $\eta$ | -0.8144 | 66.42 | 16.96 |
| 17.00 | D-dim, $\eta_{\text{eff}}$ | +0.1148 | 4.12 | **1.04** |
| 17.00 | D-dim, $\eta$ | -0.7295 | 53.21 | 13.72 |
| 20.64 | D-dim, $\eta_{\text{eff}}$ | +0.1271 | 5.06 | **1.29** |
| 20.64 | N-dim, $\eta_{\text{eff}}$ | -0.1903 | 11.04 | 2.21 |
| 20.64 | D-dim, $\eta$ | -0.5736 | 32.82 | 8.65 |
| 20.64 | N-dim, $\eta$ | -0.5461 | 29.74 | 7.87 |
| 37.63 | D-dim, $\eta_{\text{eff}}$ | +0.0935 | 2.73 | **0.67** |
| 37.63 | N-dim, $\eta_{\text{eff}}$ | -0.1320 | 5.35 | 1.11 |
| 37.63 | D-dim, $\eta$ | -0.2061 | 4.21 | 1.20 |
| 37.63 | N-dim, $\eta$ | -0.2328 | 5.37 | 1.52 |
| 55.10 | D-dim, $\eta_{\text{eff}}$ | +0.0983 | 3.02 | **0.75** |
| 55.10 | N-dim, $\eta_{\text{eff}}$ | -0.1357 | 5.64 | 1.17 |
| 55.10 | D-dim, $\eta$ | -0.2286 | 5.18 | 1.47 |
| 55.10 | N-dim, $\eta$ | -0.2179 | 4.71 | 1.34 |

*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

---

_Markdown view of https://picx.dev/p/iETSYd, served by PicX — AI-generated visual whiteboard summaries of research papers._