The reviewed record of science sign in
Pith

arxiv: 2606.05326 · v1 · pith:PTXZJ4WZ · submitted 2026-06-03 · math.OC · cs.AI· cs.LG· math-ph· math.AP· math.MP

Gradient descent at the Edge of Stability: free energy model and kinetic description of the two-layer network

Antonin Chodron de Courcel This is my paper

Reviewed by Pith T0 review T1 audit T2 compute T3 formal T4 kernel 2026-06-28 04:58 UTCgrok-4.3pith:PTXZJ4WZrecord.jsonopen to challenge →

classification math.OC cs.AIcs.LGmath-phmath.APmath.MP
keywords edge of stabilitygradient descenteffective free energymean-field limitkinetic equationWasserstein gradient flowtwo-layer networksoscillations
0
0 comments X

The pith

Gradient descent with large learning rates follows an effective free energy that includes an entropic curvature term.

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

The paper studies gradient descent when the learning rate is large enough to cause persistent oscillations in the loss and sharpness. It proposes a continuous-time model that couples the average weight trajectory with the covariance of its fast oscillations. This model shows that the relevant quantity is an effective free energy combining the risk with a curvature-related term. For wide two-layer networks, it derives a kinetic equation for the joint distribution of weights and fluctuations, interpreted as a gradient flow in Wasserstein space. Numerical tests on matrix factorization and CIFAR-10 confirm it captures oscillation envelopes.

Core claim

In the Edge of Stability regime, the dynamics of gradient descent are captured by an effective model tracking the average trajectory and time-averaged covariance of oscillations. The natural quantity to monitor is an effective free energy that augments the original risk functional with a curvature-related entropic term. For wide two-layer networks under stable oscillations, a mean-field limit yields a kinetic equation for the joint distribution of weights and fluctuations, which is the Wasserstein-2 gradient flow of a macroscopic free energy.

What carries the argument

The effective free energy, which combines the risk functional with a curvature-related entropic term, and the associated Wasserstein-2 gradient flow of the macroscopic free energy in the mean-field limit.

If this is right

  • The envelope of oscillations can be tracked even when their dynamics evolve on similar timescales as the averaged weights.
  • The model predicts spikes that occur during training of some neural network architectures.
  • For wide two-layer networks, the joint distribution of weights and fluctuations follows a novel kinetic equation.
  • This allows accurate capture of oscillation behavior in tasks like matrix factorization and deep learning on CIFAR-10.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • If the free energy model holds, it suggests that optimization at the edge of stability can be understood as minimizing a modified landscape that penalizes high curvature.
  • This kinetic description might extend to other network architectures beyond two-layer networks if similar separation of timescales applies.
  • Testing the model on deeper networks could reveal whether the mean-field limit generalizes.

Load-bearing premise

That persistent oscillations can be separated into a slow average trajectory plus fast time-averaged covariance, and that the mean-field limit applies to wide two-layer networks with stable oscillations.

What would settle it

Running gradient descent on a two-layer network with large learning rate and checking if the effective free energy decreases monotonically while the original loss oscillates.

Figures

Figures reproduced from arXiv: 2606.05326 by Antonin Chodron de Courcel.

Figure 1
Figure 1. Figure 1: Depending on the learning rate (in ascending order: blue, green, red), the trajec [PITH_FULL_IMAGE:figures/full_fig_p003_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: The effective free energy Beyond aesthetics, this observation leads us to understand that the right quantity to compare with the gradient descent evolution E( ˜θk) is not the loss E(θk) but the free energy itself (see the experiments). It is natural to analyse the stationary solution to (3.4). The important fact being that there is more stationary solutions than pure saddle points of the energy. Indeed, an… view at source ↗
Figure 3
Figure 3. Figure 3: Comparison of the trajectories of the true gradient descent dynamics and the model [PITH_FULL_IMAGE:figures/full_fig_p010_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Comparison of the trajectories of the true gradient descent dynamics and the model [PITH_FULL_IMAGE:figures/full_fig_p012_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Comparison of the trajectories of the true gradient descent dynamics and the model [PITH_FULL_IMAGE:figures/full_fig_p013_5.png] view at source ↗
read the original abstract

We study the dynamics of gradient descent in the Edge of Stability regime, where the learning rate is large enough to induce persistent oscillations in the loss and the sharpness. We propose a continuous-time effective model that tracks the evolution of the average trajectory coupled with the time-averaged covariance of its fast oscillations. Our analysis reveals that the natural quantity to monitor in such unstable regimes is an effective free energy, which combines the original risk functional with a curvature-related "entropic" term. Our model allows us to track the envelope of the oscillations even in situations where its dynamics evolve on similar timescales as the averaged weights. Otherwise stated, we can track the spikes that occur during the training of some neural network architectures. For wide two-layer neural networks optimized under stable non-vanishing oscillations, we derive a mean-field limit that results in a novel kinetic equation describing the joint distribution of weights and their fluctuations. We show that this equation can be interpreted as a Wasserstein-2 gradient flow of a macroscopic free energy. Finally, we provide numerical evidence on matrix factorization and deep learning tasks (CIFAR-10) to demonstrate the model's accuracy in capturing the envelope of the oscillations and the predictive power of the effective free energy.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

1 major / 2 minor

Summary. The manuscript studies gradient descent dynamics in the Edge of Stability regime, where large learning rates induce persistent oscillations in loss and sharpness. It proposes a continuous-time effective model coupling the slow average trajectory with the time-averaged covariance of fast oscillations, identifies an effective free energy (risk plus curvature-related entropic term) as the natural monitor, and derives a mean-field limit for wide two-layer networks yielding a novel kinetic equation for the joint law of weights and fluctuations. This PDE is shown to be the Wasserstein-2 gradient flow of a macroscopic free energy. Numerical evidence on matrix factorization and CIFAR-10 is provided to support accuracy in tracking oscillation envelopes.

Significance. If the mean-field closure and free-energy interpretation hold without requiring strict timescale separation, the work supplies a new analytic tool for oscillatory training regimes that standard averaged analyses cannot capture, with direct relevance to understanding loss spikes and effective dynamics in overparameterized networks.

major comments (1)
  1. [Mean-field limit derivation] Mean-field limit section (derivation of the kinetic equation as W2 gradient flow): the closure of the time-averaged covariance into an autonomous equation for the slow marginal is asserted to hold under 'stable non-vanishing oscillations' even when periods are comparable to the mean-weight drift. The derivation steps that justify this closure (without implicit separation assumptions) must be made fully explicit, as any residual dependence on fast-slow separation would render the resulting PDE non-autonomous and undermine the free-energy interpretation.
minor comments (2)
  1. Notation for the effective free energy functional should be introduced with an explicit equation number at first appearance to aid cross-referencing with the kinetic equation.
  2. The numerical section would benefit from an additional panel or table quantifying the error between the predicted envelope and observed spikes across multiple random seeds.

Simulated Author's Rebuttal

1 responses · 0 unresolved

We thank the referee for their careful reading and constructive feedback. We address the single major comment below.

read point-by-point responses

  1. Referee: [Mean-field limit derivation] Mean-field limit section (derivation of the kinetic equation as W2 gradient flow): the closure of the time-averaged covariance into an autonomous equation for the slow marginal is asserted to hold under 'stable non-vanishing oscillations' even when periods are comparable to the mean-weight drift. The derivation steps that justify this closure (without implicit separation assumptions) must be made fully explicit, as any residual dependence on fast-slow separation would render the resulting PDE non-autonomous and undermine the free-energy interpretation.

    Authors: We agree that the steps justifying the autonomous closure must be stated explicitly. The manuscript already asserts that the closure holds under the sole assumption of stable non-vanishing oscillations (without requiring strict timescale separation), because the time average is taken with respect to the fast periodic orbit whose shape is determined instantaneously by the current slow state. In the revision we will expand the mean-field section with a self-contained derivation that (i) defines the averaging operator along the fast orbit, (ii) shows that the resulting covariance functional depends only on the slow marginal, and (iii) verifies that the resulting kinetic equation remains autonomous and is the Wasserstein-2 gradient flow of the macroscopic free energy. No additional separation hypothesis is introduced. revision: yes

Circularity Check

0 steps flagged

No circularity: derivation chain is self-contained.

full rationale

The abstract and description present an effective continuous-time model for Edge-of-Stability oscillations, followed by a mean-field limit yielding a novel kinetic PDE interpreted as a W2 gradient flow of a macroscopic free energy. No quoted equations reduce a claimed prediction or uniqueness result to a fitted parameter, self-citation, or definitional tautology. The timescale-separation assumption is stated explicitly as a modeling choice rather than derived from prior self-work, and numerical validation on matrix factorization and CIFAR-10 is offered as independent support. The central construction therefore retains independent mathematical content.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Only abstract available; no explicit free parameters, axioms, or invented entities can be identified from the provided text.

pith-pipeline@v0.9.1-grok · 5760 in / 1150 out tokens · 25038 ms · 2026-06-28T04:58:47.409977+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

32 extracted references · 4 canonical work pages

  1. [1]

    Gradient descent on neural networks typically occurs at the edge of stability,

    Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability , author=. arXiv preprint arXiv:2103.00065 , year=

    work page arXiv

  2. [2]

    Proceedings of Machine Learning Research , year=

    Understanding Gradient Descent on the Edge of Stability in Deep Learning , author=. Proceedings of Machine Learning Research , year=

  3. [3]

    International Conference on Learning Representations , year=

    Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability , author=. International Conference on Learning Representations , year=

  4. [4]

    arXiv preprint arXiv:2210.03294 , year=

    Understanding Edge-of-Stability Training Dynamics with a Minimalist Example , author=. arXiv preprint arXiv:2210.03294 , year=

    work page arXiv

  5. [5]

    arXiv preprint arXiv:2410.24206 , year=

    Understanding Optimization in Deep Learning with Central Flows , author=. arXiv preprint arXiv:2410.24206 , year=

    work page arXiv

  6. [6]

    International Conference on Learning Representations , year=

    Sharpness-aware Minimization for Efficiently Improving Generalization , author=. International Conference on Learning Representations , year=

  7. [7]

    Proceedings of the 34th International Conference on Machine Learning (ICML) , volume =

    On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima , author =. Proceedings of the 34th International Conference on Machine Learning (ICML) , volume =. 2017 , publisher =

    2017

  8. [8]

    Advances in Neural Information Processing Systems (NeurIPS) , volume=

    On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport , author=. Advances in Neural Information Processing Systems (NeurIPS) , volume=

  9. [9]

    Proceedings of the 38th International Conference on Machine Learning , pages =

    Tensor Programs IV: Feature Learning in Infinite-Width Neural Networks , author =. Proceedings of the 38th International Conference on Machine Learning , pages =. 2021 , editor =

    2021

  10. [10]

    Tensor Programs

    Greg Yang and Dingli Yu and Chen Zhu and Soufiane Hayou , booktitle=. Tensor Programs. 2024 , url=

    2024

  11. [11]

    2025 , eprint=

    Global Convergence and Rich Feature Learning in L -Layer Infinite-Width Neural Networks under P Parametrization , author=. 2025 , eprint=

    2025

  12. [12]

    2025 , eprint=

    The Hidden Width of Deep ResNets: Tight Error Bounds and Phase Diagrams , author=. 2025 , eprint=

    2025

  13. [13]

    Advances in Neural Information Processing Systems , volume=

    On Lazy Training in Differentiable Programming , author=. Advances in Neural Information Processing Systems , volume=

  14. [14]

    Advances in Neural Information Processing Systems (NeurIPS) , volume =

    Analyzing Sharpness along GD Trajectory: Progressive Sharpening and Edge of Stability , author =. Advances in Neural Information Processing Systems (NeurIPS) , volume =. 2022 , url =

    2022

  15. [15]

    The Thirteenth International Conference on Learning Representations , year=

    Learning Dynamics of Deep Matrix Factorization Beyond the Edge of Stability , author=. The Thirteenth International Conference on Learning Representations , year=

  16. [16]

    Proceedings of the 40th International Conference on Machine Learning , pages =

    Beyond the Edge of Stability via Two-step Gradient Updates , author =. Proceedings of the 40th International Conference on Machine Learning , pages =. 2023 , editor =

    2023

  17. [17]

    Thirty-seventh Conference on Neural Information Processing Systems , year=

    What is the Inductive Bias of Flatness Regularization? A Study of Deep Matrix Factorization Models , author=. Thirty-seventh Conference on Neural Information Processing Systems , year=

  18. [18]

    Understanding Sharpness Dynamics in

    Geonhui Yoo and Minhak Song and Chulhee Yun , booktitle=. Understanding Sharpness Dynamics in. 2025 , url=

    2025

  19. [19]

    The Thirteenth International Conference on Learning Representations , year=

    Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos , author=. The Thirteenth International Conference on Learning Representations , year=

  20. [20]

    Edward J Hu and yelong shen and Phillip Wallis and Zeyuan Allen-Zhu and Yuanzhi Li and Shean Wang and Lu Wang and Weizhu Chen , booktitle=. Lo. 2022 , url=

    2022

  21. [21]

    Regis and S

    Rod Flow: A Continuous-Time Model for Gradient Descent at the Edge of Stability , author=. arXiv preprint arXiv:2602.01480 , year=

    work page arXiv

  22. [22]

    arXiv preprint , year=

    A Rod Flow Model for Adam at the Edge of Stability , author=. arXiv preprint , year=

  23. [23]

    Neural Computation , volume=

    Flat minima , author=. Neural Computation , volume=. 1997 , publisher=

    1997

  24. [24]

    1999 , publisher=

    Nonlinear Programming , author=. 1999 , publisher=

    1999

  25. [25]

    2004 , publisher=

    Introductory Lectures on Convex Optimization: A Basic Course , author=. 2004 , publisher=

    2004

  26. [26]

    Proceedings of the National Academy of Sciences , volume=

    A mean field view of the landscape of two-layer neural networks , author=. Proceedings of the National Academy of Sciences , volume=. 2018 , publisher=

    2018

  27. [27]

    Advances in Neural Information Processing Systems , volume=

    Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks , author=. Advances in Neural Information Processing Systems , volume=

  28. [28]

    Journal of Machine Learning Research , volume=

    A kernel two-sample test , author=. Journal of Machine Learning Research , volume=

  29. [29]

    Proceedings of the Thirty-First Conference on Uncertainty in Artificial Intelligence , pages=

    Training generative neural networks via Maximum Mean Discrepancy optimization , author=. Proceedings of the Thirty-First Conference on Uncertainty in Artificial Intelligence , pages=

  30. [30]

    Advances in Neural Information Processing Systems (NeurIPS) , volume=

    Maximum Mean Discrepancy Gradient Flow , author=. Advances in Neural Information Processing Systems (NeurIPS) , volume=

  31. [31]

    Advances in Neural Information Processing Systems , volume=

    Learning and Efficiency of Gradient Descent at the Edge of Stability , author=. Advances in Neural Information Processing Systems , volume=

  32. [32]

    2009 , publisher=

    Information, physics, and computation , author=. 2009 , publisher=

    2009