Why Gradients Rapidly Increase Near the End of Training

Aaron Defazio

REVIEW 3 major objections 5 minor 5 cited by

Why Gradients Rapidly Increase Near the End of Training

T0 review · 3 major / 5 minor · reviewed 2026-08-07 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read Weight decay interacting with the learning-rate schedule causes the end-of-training gradient spike in LLMs, and a corrected decay term removes it.

desk verdict A clean equilibrium explanation for end-of-training gradient blow-up on normalized layers, with a cheap fix that works on ImageNet—but the Transformer headline rests on an unverified analogy. read the letter →

arxiv 2506.02285 v2 pith:FIIRWZ7E submitted 2025-06-02 cs.LG cs.AI

Aaron Defazio This is my paper · ORCID

classification cs.LGcs.AI

keywords gradientnormweightdecaylearningrateschedulenormalizationlayerslargelanguagemodeltrainingAdamWlayerbalancingC

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper argues that the rapid rise in gradient norm seen near the end of long LLM training runs is not a mysterious instability but a predictable side effect of weight decay pulling normalized layers toward a moving equilibrium. For a layer directly followed by a normalization operation, the ratio of gradient norm to weight norm is driven to $\|g_t\|/\|x_t\| = \sqrt{2\lambda/\gamma_t}$, where $\gamma_t$ is the current learning rate, so as the schedule decays the target gradient magnitude grows. The same steady-state argument is carried over to AdamW, and uncoupled weight decay by itself does not remove the problem. The paper proposes a correction, $\hat{\lambda}_t = \lambda \gamma_t/\gamma_{\max}$, that keeps the target ratio constant, and shows that the corrected optimizers AdamC and SGDC eliminate the end-of-training spike while giving lower loss on a 120M-parameter LLM and on ImageNet ResNet-50.

What carries the argument

The central object is the ratio $\|g_t\|/\|x_t\|$ for a layer immediately followed by a normalization operation. The load-bearing identity is the orthogonality condition $\langle g_t, x_t\rangle = 0$, which simplifies the norm dynamics of the weight update to a form whose steady state is $\|g_t\|/\|x_t\| = \sqrt{2\lambda/\gamma_t}$. The paper uses this moving equilibrium both to explain the end-of-training gradient spike and to design the corrected weight-decay coefficient $\hat{\lambda}_t = \lambda \gamma_t/\gamma_{\max}$, which freezes the target ratio at a constant value.

What would settle it

Train a small transformer with a cosine learning-rate schedule, log the per-layer inner product $\langle g_t, x_t\rangle$ alongside the gradient and weight norms, and compare the observed ratio $\|g_t\|/\|x_t\|$ to the predicted $\sqrt{2\lambda/\gamma_t}$. If the ratio follows the prediction even when the cross term is non-negligible, the orthogonality assumption is harmless; if the trajectories diverge once the cross term is included, the orthogonality assumption is what breaks.

Watch

Extended reading notes

Core claim

The central discovery is that weight decay and the learning-rate schedule together set a time-varying target for the gradient-to-weight norm ratio of normalized layers. Starting from the SGD update with weight decay, the orthogonality identity $\langle g_t, x_t\rangle = 0$ for a layer followed by normalization reduces the squared-norm update to $\|x_{t+1}\|^2 = (1 - \gamma \lambda)^2 \|x_t\|^2 + \gamma^2 \|g_t\|^2$. Setting $\|x_{t+1}\| = \|x_t\|$ at the steady state and dropping the small $\lambda^2$ term yields $\|g_t\|/\|x_t\| = \sqrt{2\lambda/\gamma_t}$. When $\gamma_t$ decays toward zero under a cosine or similar schedule, this target ratio blows up, which is the mechanism behind the observed gradient-norm increase. For AdamW the same argument holds in a weighted norm, and the paper shows the original Adam's coupled weight decay does not admit a uniform ratio, which it presents as an explanation for AdamW's empirical advantage. The proposed corrected decay, $\hat{\lambda}_t = \lambda \gamma_t/\gamma_{\max}$, fixes the time dependence and, according to the experiments, removes the gradient blow-up while improving loss.

Load-bearing premise

The argument rests on assuming the gradient and the weights of each normalized layer are exactly orthogonal, an identity that holds for a layer directly followed by a normalization layer but is only approximate for transformer layers, where the paper nevertheless applies the correction to essentially all linear layers.

Editorial extensions

If this is right

On the 120M-parameter LLM trained for 200B tokens, AdamC removes the second-half gradient-norm increase and yields lower loss than AdamW throughout training.
On ImageNet ResNet-50, SGDC eliminates the rapid end-of-training gradient spike that SGDM exhibits, while slightly improving test accuracy (77.07% vs 76.95%).
Because the corrected weight decay keeps the gradient-to-weight ratio constant, it also stabilizes weight norms, which otherwise decay rapidly under a cosine schedule.
The analysis implies that AdamW's decoupled weight decay, not just its adaptive steps, is what allows all normalized layers to settle at the same ratio; Adam's coupled decay leaves layers unbalanced.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

If the ratio $\|g_t\|/\|x_t\|$ is the right control parameter, monitoring it during training could give an early signal that the schedule is pushing the model into the tail-blow-up regime before the loss degrades.
The same ratio-balancing logic suggests the corrected weight decay might also make the optimal weight-decay coefficient less sensitive to schedule length and warmup, since it removes the schedule's time dependence from the steady-state target.
The paper applies the correction to all linear layers except the output; a more precise rule based on whether a layer's activations are normalized, or on the measured cross term $\langle g_t, x_t\rangle$, could be tested to see whether it further improves transformer training.
A natural stress test is to apply AdamC to much longer runs or to schedules with more aggressive decay, where the uncorrected spike is largest; the correction should keep the gradient norm flat if the mechanism is complete.

Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

Short version: the paper gives a clean steady-state explanation for why gradient norms rise at the end of training under cosine schedules, and proposes a one-line fix: scale weight decay by γt/γmax. The ImageNet ResNet experiments genuinely support the mechanism. The LLM part, which is the headline, is where it gets soft: the theory assumes ⟨g, x⟩ = 0, which holds for a layer directly followed by normalization, but Llama's residual structure doesn't satisfy that, and the paper admits it. So the paper doesn't actually derive the Transformer tail blow-up; it asserts it by analogy and then shows the fix still works empirically.

What's new: the ratio interpretation and the specific corrected decay. The underlying equilibrium is Van Laarhoven's, and Xie et al. already suggested scheduling weight decay, but I don't think they gave this particular constructive form or connected it to the observed gradient-norm shape. The math in Section 3 is correct under the stated assumptions, and the figures showing layer ratios tracking sqrt(2λ/γ) are convincing for ResNet.

Soft spots, in order. First, the Transformer extrapolation is the load-bearing claim. The cross term ⟨g, x⟩ is not zero for attention and MLP projections; measuring it directly would decide whether Eq. (2) governs the LLM tail or whether AdamC just reduces late-training weight decay. The paper's own Section 6 concession doesn't help; it applies the correction to all linear layers anyway. Second, the "lower loss" evidence is thin: the ImageNet gain is ~0.1% and the LLM gain is 0.004 loss, with no error bars on the LLM run. A matched baseline controlling total weight decay would be needed to show the correction, not just less decay, is responsible. Third, novelty is moderate, but the paper is honest about prior work.

Who is this for? People training long LLM runs and optimizer theorists. It's a useful note, not a breakthrough. I'd send it out for review; the mechanism is plausible, the fix is cheap, and a referee can ask for the direct measurement that would settle the Transformer question.

Referee Report

3 major / 5 minor

Summary. This short note addresses the rapid increase of gradient norms near the end of long LLM training runs. The paper models a normalized linear layer under weight decay and learning-rate scheduling, deriving the steady-state ratio ||g_t||/||x_t|| = sqrt(2λ/γ_t) (Section 3, Eq. 2). As γ_t decays, this target ratio grows, which the paper identifies as the cause of the tail blow-up. The proposed fix is a corrected weight decay λ̂_t = λ γ_t/γ_max applied only to normalized layers, yielding the AdamC and SGDC optimizers. Experiments include a 200B-token 120M-parameter Llama-3 run on FineWeb-Edu and ImageNet ResNet-50 training; the paper reports that the correction eliminates the late-training gradient increase and gives lower loss.

Significance. If valid, the paper offers a simple, parameter-free explanation and a practical optimizer correction for a phenomenon observed in real LLM training. The theoretical steady-state prediction in Eq. (2) is falsifiable and is convincingly validated on ImageNet ResNet-50 with a flat schedule (Fig. 2) and a cosine schedule (Fig. 3), and the five-seed ImageNet accuracy comparison (Fig. 5) is a genuine strength. The main weaknesses are the extrapolation of the theory to Transformers—where the orthogonality premise does not hold—and the thin single-run LLM evidence for the lower-loss claim. The paper is clearly written and the algorithm is straightforward to implement, but these gaps prevent the headline LLM claim from being fully established.

major comments (3)

[Section 6, Figure 4] The steady-state ratio (Eq. 2) is derived from Eq. (1), whose second equality drops the cross term ⟨g_t, x_t⟩ under the premise that the layer is immediately followed by a normalization operation. In the Llama-style architecture used in Figure 4, no linear layer has this structure: the QKV projections are preceded by RMSNorm, and the attention/MLP output projections feed the residual stream before any normalization. The text itself concedes this ('The application of normalization operations in Transformers no longer corresponds to the direct linear-followed-by-norm structure that our theory covers'), yet the correction is applied to essentially every linear layer. Because the cross term is not guaranteed to vanish, the observed LLM gradient-norm behavior and the success of AdamC cannot be attributed to the proposed mechanism without a direct measurement of ⟨g_t, x_t⟩ (or of the residual corrections) on at least one representative Transformer layer. This is the central load-bearing gap for the LLM claim.
[Section 5] The uncoupled update x_{t+1} = x_t - γ_t g_t - λ_t x_t gives the steady-state ratio ∥g_t∥/∥x_t∥ = √(2λ_t)/γ_t. Substituting λ̂_t = λ γ_t/γ_max yields √(2λ γ_t/γ_max)/γ_t = √(2λ/(γ_max γ_t)), not √(2λ/γ_max) as claimed. The proposed correction makes the ratio schedule-independent only when the weight-decay term is multiplied by the learning rate, as in Algorithm 1's line 12, where the decay is γ_t^2 λ/γ_max. The manuscript should correct this derivation, define SGDC's update precisely (coupled or uncoupled), and, if the uncoupled form is intended, use λ γ_t^2/γ_max as the corrected decay.
[Figure 4] The 'lower loss throughout training' claim for LLMs rests on a single 200B-token run with no error bars or seed variation. The final loss difference (2.461 vs 2.457) is within plausible run-to-run variation. Given that the Transformer extension is already heuristic, the loss claim needs at least three independent runs or a checkpoint-level uncertainty estimate.

minor comments (5)

[Section 5] The first display equation in Section 5 uses an undefined symbol η; this is likely a typo for γ_t and should be fixed.
[Algorithm 1] Line 12's decay term γ_t^2 λ/γ_max should be explicitly linked to λ̂_t = λ γ_t/γ_max in the text, since the appearance of γ_t^2 is otherwise surprising.
[Section 2] The phrase 'it's relation' should be 'its relation'.
[Figures 2 and 3] The y-axis labels appear as '10 2' and '10 1' due to missing superscripts; please format them as 10^{-2} and 10^{-1}.
[Section 4.2] The Hölder-inequality approximation for ∥x_t∥_{A_t} is explicitly non-rigorous; a sentence noting that the infinity-norm interpretation is heuristic would help prevent over-reading.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity: the gradient-to-weight ratio is a closed-form steady-state prediction and the AdamC correction is validated out-of-sample.

full rationale

I find no significant circularity in this paper. The central relation ||g_t||/||x_t|| = sqrt(2*lambda/gamma_t) is derived algebraically from a norm-update identity and the orthogonality property of normalized layers, and it is then compared against measured smoothed gradient-to-weight ratios in Figures 2 and 3 with no fitted constants. The corrected weight decay, lambda_hat_t = lambda * gamma_t / gamma_max, is obtained by substituting into that same closed-form expression; the paper does not fit this correction to the gradient-norm data it later reports. The claims that AdamC and SGDC remove the tail increase in gradient norm and improve loss are tested on external benchmarks (ImageNet ResNet-50 and a 120M-parameter LLM on FineWeb-Edu), so the outcome is not forced by construction. The self-citations used in the paper, such as Defazio and Gower (2021) for the effective learning rate under momentum and Defazio and Bottou (2022) for layer balancing, are supporting and independently checkable; they are not the load-bearing source of the main prediction. The paper itself explicitly flags a real scope limitation in Section 6: "The application of normalization operations in Transformers no longer corresponds to the direct linear-followed-by-norm structure that our theory covers," and yet the correction is applied to essentially all linear layers. That is a correctness or generalization risk for the LLM extrapolation, not a circularity, because the LLM experiments still constitute an out-of-sample empirical check rather than a restatement of the theory's assumptions. No derivation step reduces by definition to its inputs, and no fitted parameter is renamed as a prediction. Therefore the appropriate circularity score is 0, with the noted Transformer applicability caveat belonging to the correctness assessment rather than to circularity.

Assumptions & free parameters 0 free parameters · 5 assumptions · 0 invented entities

The central prediction rests on a small number of structural assumptions about normalization and steady-state behavior; the paper states most of them. No free parameters are fit, but the experimental validation tunes the learning rate per optimizer, and the Transformer application relies on an extension the theory does not strictly cover.

assumptions (5)

domain assumption For a layer immediately followed by normalization, the gradient is orthogonal to the weights, <g_t, x_t> = 0.
Used in Section 3, Eq. (1), to drop the cross term in the norm recurrence. True for scale-invariant layers, but only approximately for residual and Transformer structures.
domain assumption The weight norm dynamics reach a steady state satisfying ||x_{t+1}|| = ||x_t||.
Section 3 solves for this equilibrium and the paper acknowledges it is an assumption for a noisy, non-stationary system.
standard math The second-order term lambda^2 * gamma^2 is negligible.
Section 3 drops this term to obtain the simplified equilibrium ratio sqrt(2*lambda/gamma).
domain assumption For AdamW, momentum is ignored, A_t ~ A_{t+1}, and the weighted-norm steady-state analysis carries over.
Section 4 states that the AdamW argument is approximate and assumes no momentum, so the derived ratio is a first-order approximation.
ad hoc to paper The weighted norm approximations ||g||_{A^{-1}} ~ sqrt(||g||_1) and ||x||_A ~ ||x||_inf * sqrt(||g||_1) hold after the SignSGD approximation.
Section 4.2 uses these non-rigorous approximations to derive the infinity-norm balance interpretation of AdamW.

how reviews work

0 comments

Cite this review

Pith. "Pith review of Why Gradients Rapidly Increase Near the End of Training." pith.science (2026). https://pith.science/paper/FIIRWZ7E

@misc{pith2026250602285,
  author       = {Pith},
  title        = {Pith review of: Why Gradients Rapidly Increase Near the End of Training},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/FIIRWZ7E}},
  note         = {Machine review of arXiv:2506.02285}
}

read the original abstract

During long-duration Large Language Model (LLM) training runs the gradient norm increases rapidly near the end of training. In this short note, we show that this increase is due to an unintended interaction between weight decay, normalization layers, and the learning rate schedule. We propose a simple correction that fixes this behavior while also resulting in lower loss values throughout training.

Discussion (0). Continue with ORCID to comment.

Forward citations

Cited by 5 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

Safeguarded Stochastic Polyak Step Sizes for Non-smooth Optimization: Robust Performance Without Small (Sub)Gradients
math.OC 2025-12 conditional novelty 7.0 of 10

A safeguarded stochastic Polyak step size, SPS_safe, yields O(1/√T) convergence to a neighborhood for convex non-smooth problems without interpolation or oracle loss values, with a momentum variant.
Full-bandwidth transformer
cs.AI 2026-08 conditional novelty 6.0 of 10

Fusing the previous top-layer hidden state into the next input token through a gated linear unit improves transformer pretraining data efficiency and decoding accuracy at 1B scale.
OmniOpt: Taxonomy, Geometry, and Benchmarking of Modern Optimizers
cs.LG 2026-07 conditional novelty 6.0 of 10

A meta-pipeline plus LMO four-axis view yields a dual taxonomy of 108 optimizers, and a multi-objective LLM/vision benchmark shows no single family dominates the quality–cost–memory frontier.
Scale Weight Decay and Train Better
cs.LG 2026-07 conditional novelty 5.0 of 10

Muon with weight decay scaled by η/η_max reaches the same MoE validation loss ~30% faster than constant-decay Muon while preserving asymptotic stationarity of the unregularized objective.
Why Do We Need Warm-up? A Theoretical Perspective
cs.LG 2025-10 conditional novelty 5.0 of 10

Under the proposed (H0,H1)-smoothness condition, gradient descent with a warm-up-style adaptive step-size provably converges faster than with any fixed step-size.

Reference graph

Works this paper leans on

22 extracted references · 14 canonical work pages · cited by 5 Pith papers

[1]

L., Kiros, J

Ba, J. L., Kiros, J. R., and Hinton, G. E. (2016). Layer normalization. arXiv preprint arXiv:1607.06450

arXiv 2016
[2]

Bernstein, J., Vahdat, A., Yue, Y., and Liu, M.-Y. (2020). On the distance between two neural networks and the stability of learning. Advances in Neural Information Processing Systems , 33:21370--21381

work page 2020
[3]

Bernstein, J., Wang, Y.-X., Azizzadenesheli, K., and Anandkumar, A. (2018). signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning , pages 560--569. PMLR

work page 2018
[4]

D'Angelo, F., Andriushchenko, M., Varre, A., and Flammarion, N. (2024). Why do we need weight decay in modern deep learning?

work page 2024
[5]

and Bottou, L

Defazio, A. and Bottou, L. (2022). A scaling calculus for the design and initialization of relu networks. Neural Comput. Appl. , 34(17):14807–14821

work page 2022
[6]

Defazio, A., Cutkosky, A., Mehta, H., and Mishchenko, K. (2023). When, why and how much? adaptive learning rate scheduling by refinement

work page 2023
[7]

and Gower, R

Defazio, A. and Gower, R. M. (2021). The power of factorial powers: New parameter settings for (stochastic) optimization. In Balasubramanian, V. N. and Tsang, I., editors, Proceedings of The 13th Asian Conference on Machine Learning , volume 157 of Proceedings of Machine Learning Research , pages 49--64. PMLR

2021
[8]

and Mishchenko, K

Defazio, A. and Mishchenko, K. (2023). Learning-rate-free learning by D -adaptation. The 40th International Conference on Machine Learning (ICML 2023)

2023

Show all 22 references

[9]

Duchi, J., Hazan, E., and Singer, Y. (2011). Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research , 12(61)

2011
[10]

He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition

2016
[11]

and Szegedy, C

Ioffe, S. and Szegedy, C. (2015). Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning , pages 448--456. pmlr

2015
[12]

Kosson, A., Messmer, B., and Jaggi, M. (2024). Rotational equilibrium: How weight decay balances learning across neural networks

2024
[13]

and Hutter, F

Loshchilov, I. and Hutter, F. (2019). Decoupled weight decay regularization. In International Conference on Learning Representations

2019
[14]

Orabona, F. (2019). A modern introduction to online learning. arXiv preprint arXiv:1912.13213

2019 arXiv
[15]

B., Lozhkov, A., Mitchell, M., Raffel, C., Werra, L

Penedo, G., Kydlíček, H., allal, L. B., Lozhkov, A., Mitchell, M., Raffel, C., Werra, L. V., and Wolf, T. (2024). The fineweb datasets: Decanting the web for the finest text data at scale

2024
[16]

C., and Fei-Fei, L

Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A. C., and Fei-Fei, L. (2015). ImageNet Large Scale Visual Recognition Challenge . International Journal of Computer Vision (IJCV) , 115(3)

2015
[17]

Schaipp, F. (2023). Decay no more. In ICLR Blogposts 2023 . https://iclr-blogposts.github.io/2023/blog/2023/adamw/

2023
[18]

Van Laarhoven, T. (2017). L2 regularization versus batch and weight normalization. arXiv preprint arXiv:1706.05350

2017 arXiv
[19]

and Li, Z

Xie, S. and Li, Z. (2024). Implicit bias of adamw: _ norm constrained optimization. arXiv preprint arXiv:2404.04454

2024 arXiv
[20]

Xie, Z., zhiqiang xu, Zhang, J., Sato, I., and Sugiyama, M. (2023). On the overlooked pitfalls of weight decay and how to mitigate them: A gradient-norm perspective. In Thirty-seventh Conference on Neural Information Processing Systems

2023
[21]

Zhang, G., Wang, C., Xu, B., and Grosse, R. (2019). Three mechanisms of weight decay regularization. In International Conference on Learning Representations

2019
[22]

Zhuang, Z., Liu, M., Cutkosky, A., and Orabona, F. (2022). Understanding adamw through proximal methods and scale-freeness. Transactions on Machine Learning Research . Expert Certification

2022

Pith tools

Reviewed August 7, 2026 · model on record in the stance chip above.

[1] [1]

L., Kiros, J

Ba, J. L., Kiros, J. R., and Hinton, G. E. (2016). Layer normalization. arXiv preprint arXiv:1607.06450

arXiv 2016

[2] [2]

Bernstein, J., Vahdat, A., Yue, Y., and Liu, M.-Y. (2020). On the distance between two neural networks and the stability of learning. Advances in Neural Information Processing Systems , 33:21370--21381

work page 2020

[3] [3]

Bernstein, J., Wang, Y.-X., Azizzadenesheli, K., and Anandkumar, A. (2018). signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning , pages 560--569. PMLR

work page 2018

[4] [4]

D'Angelo, F., Andriushchenko, M., Varre, A., and Flammarion, N. (2024). Why do we need weight decay in modern deep learning?

work page 2024

[5] [5]

and Bottou, L

Defazio, A. and Bottou, L. (2022). A scaling calculus for the design and initialization of relu networks. Neural Comput. Appl. , 34(17):14807–14821

work page 2022

[6] [6]

Defazio, A., Cutkosky, A., Mehta, H., and Mishchenko, K. (2023). When, why and how much? adaptive learning rate scheduling by refinement

work page 2023

[7] [7]

and Gower, R

Defazio, A. and Gower, R. M. (2021). The power of factorial powers: New parameter settings for (stochastic) optimization. In Balasubramanian, V. N. and Tsang, I., editors, Proceedings of The 13th Asian Conference on Machine Learning , volume 157 of Proceedings of Machine Learning Research , pages 49--64. PMLR

2021

[8] [8]

and Mishchenko, K

Defazio, A. and Mishchenko, K. (2023). Learning-rate-free learning by D -adaptation. The 40th International Conference on Machine Learning (ICML 2023)

2023