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

• URL: http://arxiv.org/abs/2606.05326v1
• Date: Wed, 03 Jun 2026 18:11:33 GMT
• Title: Gradient descent at the Edge of Stability: free energy model and kinetic description of the two-layer network
• Abstract summary: We study the dynamics of gradient descent in the Edge of Stability regime.<n>We propose a continuous-time effective model that tracks the evolution of the average trajectory.<n>We show that this equation can be interpreted as a Wasserstein-2 gradient flow of a macroscopic free energy.
• Score: 0.0
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• 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.

Related papers

• Detecting Deepfakes via Hamiltonian Dynamics [58.881686693914446]
  We propose a new perspective on deepfake detection: moving from static pattern recognition to dynamical stability analysis.<n>Our approach is motivated by physics-inspired priors.<n>Hamiltonian Action Anomaly Detection (HAAD) outperforms evaluated state-of-the-art baselines.
  arXiv  Detail & Related papers  (2026-05-06T01:55:38Z)
• Stochastic Gradient Descent in the Saddle-to-Saddle Regime of Deep Linear Networks [74.46751089984072]
  gradient descent (SGD) noise on descent in deep linear networks is poorly understood.<n>We investigate the dynamics of SGD during training of DLNs in the saddle-to-saddle regime.<n>Our results establish that SGD noise encodes information about the progression of feature learning but does not fundamentally alter the saddle-to-saddle dynamics.
  arXiv  Detail & Related papers  (2026-04-07T18:43:08Z)
• Non-Equilibrium Stochastic Dynamics as a Unified Framework for Insight and Repetitive Learning: A Kramers Escape Approach to Continual Learning [0.0]
  Continual learning in artificial neural networks is fundamentally limited by the stability-plasticity dilemma.<n>We show that both problems admit a common resolution within non-equilibrium statistical physics.
  arXiv  Detail & Related papers  (2026-04-05T15:42:23Z)
• Identifiable learning of dissipative dynamics [25.409059056398124]
  We introduce I-OnsagerNet, a neural framework that learns dissipative dynamics directly from trajectories.<n>I-OnsagerNet extends the Onsager principle to guarantee that the learned potential is obtained from the stationary density.<n>Our approach enables us to calculate the entropy production and to quantify irreversibility, offering a principled way to detect and quantify deviations from equilibrium.
  arXiv  Detail & Related papers  (2025-10-28T07:57:14Z)
• Langevin Flows for Modeling Neural Latent Dynamics [81.81271685018284]
  We introduce LangevinFlow, a sequential Variational Auto-Encoder where the time evolution of latent variables is governed by the underdamped Langevin equation.<n>Our approach incorporates physical priors -- such as inertia, damping, a learned potential function, and forces -- to represent both autonomous and non-autonomous processes in neural systems.<n>Our method outperforms state-of-the-art baselines on synthetic neural populations generated by a Lorenz attractor.
  arXiv  Detail & Related papers  (2025-07-15T17:57:48Z)
• Dynamic metastability in the self-attention model [22.689695473655906]
  We consider the self-attention model - an interacting particle system on the unit sphere - which serves as a toy model for Transformers. We prove the appearance of dynamic metastability conjectured in [GLPR23]. We show that under an appropriate time-rescaling, the energy reaches its global maximum in finite time and has a staircase profile.
  arXiv  Detail & Related papers  (2024-10-09T12:50:50Z)
• An optimization-based equilibrium measure describes non-equilibrium steady state dynamics: application to edge of chaos [2.5690340428649328]
  Understanding neural dynamics is a central topic in machine learning, non-linear physics and neuroscience. The dynamics is non-linear, and particularly non-gradient, i.e., the driving force can not be written as gradient of a potential.
  arXiv  Detail & Related papers  (2024-01-18T14:25:32Z)
• Slow semiclassical dynamics of a two-dimensional Hubbard model in disorder-free potentials [77.34726150561087]
  We show that introduction of harmonic and spin-dependent linear potentials sufficiently validates fTWA for longer times. In particular, we focus on a finite two-dimensional system and show that at intermediate linear potential strength, the addition of a harmonic potential and spin dependence of the tilt, results in subdiffusive dynamics.
  arXiv  Detail & Related papers  (2022-10-03T16:51:25Z)
• The Limiting Dynamics of SGD: Modified Loss, Phase Space Oscillations, and Anomalous Diffusion [29.489737359897312]
  We study the limiting dynamics of deep neural networks trained with gradient descent (SGD) We show that the key ingredient driving these dynamics is not the original training loss, but rather the combination of a modified loss, which implicitly regularizes the velocity and probability currents, which cause oscillations in phase space.
  arXiv  Detail & Related papers  (2021-07-19T20:18:57Z)
• Assessment of weak-coupling approximations on a driven two-level system under dissipation [58.720142291102135]
  We study a driven qubit through the numerically exact and non-perturbative method known as the Liouville-von equation with dissipation. We propose a metric that may be used in experiments to map the regime of validity of the Lindblad equation in predicting the steady state of the driven qubit.
  arXiv  Detail & Related papers  (2020-11-11T22:45:57Z)
• Feedback-induced instabilities and dynamics in the Jaynes-Cummings model [62.997667081978825]
  We investigate the coherence and steady-state properties of the Jaynes-Cummings model subjected to time-delayed coherent feedback. The introduced feedback qualitatively modifies the dynamical response and steady-state quantum properties of the system.
  arXiv  Detail & Related papers  (2020-06-20T10:07:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.