Gradient Descent Dynamics

Ask about this topic

Overview

This field studies how parameters and predictions change during training, and how learning rates, initialization, data structure, and optimizer choice affect convergence and generalization. Researchers use gradient-flow and ODE approximations, discrete dynamical-systems tools such as Lyapunov and bifurcation analysis, and stochastic or spectral methods to study settings from linear regression and matrix factorization to neural networks. A major theme is the edge of stability, where large steps can produce oscillations that alter sharpness, feature learning, and implicit bias; related work examines early stopping, grokking, and the effects of minibatch noise and momentum. An open tension is that continuous-time and simplified-model theories often give clear guarantees, while finite-step training in modern networks can show cycles, transient amplification, and optimizer-specific behavior that those approximations miss.

Latest papers

CardsList

157 papers23 papers in the last four weeks, up 109% on the four weeks before. 0.1% of all new papers.

Apr 8, 2026

Conservation Law Breaking at the Edge of Stability: A Spectral Theory of Non-Convex Neural Network Optimization

Why does gradient descent reliably find good solutions in non-convex neural network optimization, despite the landscape being NP-hard in the worst case? We show that gradient flow on L-layer ReLU networks without bias preserves L-1 conservation laws C_l = ||W_{l+1}||_F^2 - ||W_l||F^2, confining trajectories to lower-dimensional manifolds. Under discrete gradient descent, these laws break with total drift scaling as eta^alpha where alpha is approximately 1.1-1.6 depending on architecture, loss function, and width. We decompose this drift exactly as eta^2 * S(eta), where the gradient imbalance sum S(eta) admits a closed-form spectral crossover formula with mode coefficients c_k proportional to e_k(0)^2 * lambda{x,k}^2, derived from first principles and validated for both linear (R=0.85) and ReLU (R>0.80) networks. For cross-entropy loss, softmax probability concentration drives exponential Hessian spectral compression with timescale tau = Theta(1/eta) independent of training set size, explaining why cross-entropy self-regularizes the drift exponent near alpha=1.0. We identify two dynamical regimes separated by a width-dependent transition: a perturbative sub-Edge-of-Stability regime where the spectral formula applies, and a non-perturbative regime with extensive mode coupling. All predictions are validated across 23 experiments.

Daniel Nobrega Medeiros

University of Colorado Boulder

Neural Network OptimizationGradient Descent DynamicsEdge of StabilityNeural Network Stability
AskPDF
Apr 7, 2026

Stochastic Gradient Descent in the Saddle-to-Saddle Regime of Deep Linear Networks

Deep linear networks (DLNs) are used as an analytically tractable model of the training dynamics of deep neural networks. While gradient descent in DLNs is known to exhibit saddle-to-saddle dynamics, the impact of stochastic gradient descent (SGD) noise on this regime remains poorly understood. We investigate the dynamics of SGD during training of DLNs in the saddle-to-saddle regime. We model the training dynamics as stochastic Langevin dynamics with anisotropic, state-dependent noise. Under the assumption of aligned and balanced weights, we derive an exact decomposition of the dynamics into a system of one-dimensional per-mode stochastic differential equations. This establishes that the maximal diffusion along a mode precedes the corresponding feature being completely learned. We also derive the stationary distribution of SGD for each mode: in the absence of label noise, its marginal distribution along specific features coincides with the stationary distribution of gradient flow, while in the presence of label noise it approximates a Boltzmann distribution. Finally, we confirm experimentally that the theoretical results hold qualitatively even without aligned or balanced weights. These results establish that SGD noise encodes information about the progression of feature learning but does not fundamentally alter the saddle-to-saddle dynamics.

Guillaume Corlouer, Avi Semler, Alexander Strang +1

Moirai · University of Oxford · University of California, Berkeley +1

Gradient Descent DynamicsStochastic Gradient DescentDeep Linear NetworksRepresentation Learning
AskPDF
Apr 6, 2026

Grokking as Dimensional Phase Transition in Neural Networks

Neural network grokking -- the abrupt memorization-to-generalization transition -- challenges our understanding of learning dynamics. Through finite-size scaling of gradient avalanche dynamics across eight model scales, we find that grokking is a dimensional phase transition: effective dimensionality~DD crosses from sub-diffusive (subcritical, D<1D < 1) to super-diffusive (supercritical, D>1D > 1) at generalization onset, exhibiting self-organized criticality (SOC). Crucially, DD reflects gradient field geometry, not network architecture: synthetic i.i.d.\ Gaussian gradients maintain D≈1D \approx 1 regardless of graph topology, while real training exhibits dimensional excess from backpropagation correlations. The grokking-localized D(t)D(t) crossing -- robust across topologies -- offers new insight into the trainability of overparameterized networks.

Ping Wang

Institute of High Energy Physics, Chinese Academy of Science, 100049 Beijing, China

Phase Transitions in Neural NetworksGrokkingNeural Network Training DynamicsStatistical Physics of LearningGradient Descent Dynamics
AskPDF
Apr 3, 2026

Product-Stability: Provable Convergence for Gradient Descent on the Edge of Stability

Empirically, modern deep learning training often occurs at the Edge of Stability (EoS), where the sharpness of the loss exceeds the threshold below which classical convergence analysis applies. Despite recent progress, existing theoretical explanations of EoS either rely on restrictive assumptions or focus on specific squared-loss-type objectives. In this work, we introduce and study a structural property of loss functions that we term product-stability. We show that for losses with product-stable minima, gradient descent applied to objectives of the form (x,y)↦l(xy)(x,y) \mapsto l(xy) can provably converge to the local minimum even when training in the EoS regime. This framework substantially generalizes prior results and applies to a broad class of losses, including binary cross entropy. Using bifurcation diagrams, we characterize the resulting training dynamics, explain the emergence of stable oscillations, and precisely quantify the sharpness at convergence. Together, our results offer a principled explanation for stable EoS training for a wider class of loss functions.

Eric Gan

Independent Researcher.

Deep Learning OptimizationEdge of StabilityGradient Descent DynamicsNeural Network StabilityLoss Landscape Sharpness
AskPDF
Apr 2, 2026

Plateaus, Optima, and Overfitting in Multi-Layer Perceptrons: A Saddle-Saddle-Attractor Scenario

Vanishing gradients and overfitting are central problems in machine learning, yet are typically analyzed in asymptotic regimes that obscure their dynamical origins. Here we provide a dynamical description of learning in multi-layer perceptrons (MLPs) via a minimal model inspired by Fukumizu and Amari. We show that training dynamics traverse plateau and near-optimal regions, both organized by saddle structures, before converging to an overfitting regime. Under suitable conditions on the data, this regime collapses to a single attractor modulo symmetry. Furthermore, for finite noisy datasets, convergence to the theoretical optimum is impossible, and the dynamics necessarily settle into an overfitting solution.

Alex Alì Maleknia, Yuzuru Sato

Department of Mathematics, Hokkaido University, Kita 12 Nishi 7, Kita-ku, Sapporo, Hokkaido 060-0812, Japan · Present address: IMAG, IROKO, Univ Montpellier, Inria, CNRS · RIES, Hokkaido University, Kita 12 Nishi 7, Kita-ku, Sapporo, Hokkaido 060-0812, Japan

Neural Network Training DynamicsMultilayer PerceptronsNeural Network GeneralizationGradient Descent Dynamics
AskPDF
Apr 1, 2026

GRASP: Gradient Realignment via Active Shared Perception for Multi-Agent Collaborative Optimization

Non-stationarity arises from concurrent policy updates and leads to persistent environmental fluctuations. Existing approaches like Centralized Training with Decentralized Execution (CTDE) and sequential update schemes mitigate this issue. However, since the perception of the policies of other agents remains dependent on sampling environmental interaction data, the agent essentially operates in a passive perception state. This inevitably triggers equilibrium oscillations and significantly slows the convergence speed of the system. To address this issue, we propose Gradient Realignment via Active Shared Perception (GRASP), a novel framework that defines generalized Bellman equilibrium as a stable objective for policy evolution. The core mechanism of GRASP involves utilizing the independent gradients of agents to derive a defined consensus gradient, enabling agents to actively perceive policy updates and optimize team collaboration. Theoretically, we leverage the Kakutani Fixed-Point Theorem to prove that the consensus direction u∗u^* guarantees the existence and attainability of this equilibrium. Extensive experiments on StarCraft II Multi-Agent Challenge (SMAC) and Google Research Football (GRF) demonstrate the scalability and promising performance of the framework.

Sihan Zhou, Tiantian He, Yifan Lu +2

College of Computer Science and Technology, and Key Laboratory of Social Computing and Cognitive Intelligence (Ministry of Education), Dalian University of Technology, China · Center for Frontier AI Research, Institute of High Performance Computing, Singapore Institute of Manufacturing Technology, Agency for Science, Technology and Research (A*STAR), Singapore · School of Computer Science and Engineering, Nanyang Technological University, Singapore

Cooperative MARLNon-Stationary RLMulti-Agent ConsensusGradient Descent Dynamics
AskPDF
Mar 30, 2026

Immediate Derivatives Suffice for Online Recurrent Adaptation

For three decades online recurrent learning has been assumed to require propagating a Jacobian tensor through the network's dynamics at O(n4)O(n^4) per step. We show it doesn't. Dropping the propagation entirely (d=0d=0, O(n2)O(n^2) memory) matches full RTRL within CI on held-out BCI cross-session drift (TOST equivalent within ±3\pm 3 pp at n=20n=20, Adam, float64), and across vanilla-RNN synthetic cells (sine and Lorenz under Adam and SGD) and LSTM/sine under Adam. A decomposition gRTRL=gimm+gpastg_{RTRL} = g_{imm} + g_{past} explains why. On BCI, gpastg_{past} concentrates in a single direction (top-1 singular fraction 0.62-0.74 across four optimizers, vs 0.333 for gimmg_{imm}), and the four-optimizer full-RTRL-vs-d=0d=0 recovery gap tracks each optimizer's per-layer update-magnitude ratio ∥ΔWhh∥/∥ΔWout∥\|ΔW_{hh}\|/\|ΔW_{out}\| monotonically. A stationary (no-drift) control collapses both concentrations to ~0.6: the drift-specific signal is the differential, not gpastg_{past}'s absolute rank-1 structure. The signature and the behavioral gap both collapse on LSTM, consistent with a mechanism specific to additive linear recurrence. On synthetic sine, gimmg_{imm} is redundant with gpastg_{past}, which predicts the synthetic null. Full RTRL's one robust advantage is LARS (+17 to +27 pp), but d=0d=0+LARS also fails to adapt independently; the gap is an optimizer×\timesmethod interaction, not a method-quality claim. We characterize the regime: d=0d=0+Adam+float64 is robust; SGD, Adafactor, and float32 have specific fragilities documented in the paper. On the evaluated cells, the 1000×1000\times memory saving at n=1024n=1024 (O(n2)O(n^2) vs O(n4)O(n^4)) comes with no measured recovery cost.

Aur Shalev Merin

Independent Researcher

Recurrent Neural NetworksGradient Descent Dynamics
AskPDF
Mar 15, 2026

Rigorous Asymptotics for First-Order Algorithms Through the Dynamical Cavity Method

Dynamical Mean Field Theory (DMFT) provides an asymptotic description of the dynamics of macroscopic observables in certain disordered systems. Originally pioneered in the context of spin glasses by Sompolinsky and Zippelius (1982), it has since been used to derive asymptotic dynamical equations for a wide range of models in physics, high-dimensional statistics and machine learning. One of the main tools used by physicists to obtain these equations is the dynamical cavity method, which has remained largely non-rigorous. In contrast, existing mathematical formalizations have relied on alternative approaches, including Gaussian conditioning, large deviations over paths, or Fourier analysis. In this work, we formalize the dynamical cavity method and use it to give a new proof of the DMFT equations for General First Order Methods, a broad class of dynamics encompassing algorithms such as Gradient Descent and Approximate Message Passing.

Yatin Dandi, David Gamarnik, Francisco Pernice +1

Statistical Physics of Computation Laboratory, École polytechnique fédérale de Lausanne (EPFL) CH-1015 Lausanne. · Sloan School of Management, Operations Research Center and Institute of Data, Systems and Society (IDSS), MIT. · CSAIL and LIDS, MIT.

Approximate Message PassingDynamical Mean-Field TheoryMean-Field TheoryGradient Descent Dynamics
AskPDF
Mar 14, 2026

Spectral Edge Dynamics of Training Trajectories: Signal--Noise Geometry Across Scales

Despite hundreds of millions of parameters, transformer training trajectories evolve within only a few coherent directions. We introduce Spectral Edge Dynamics (SED) to quantify this structure: a rolling-window SVD of parameter updates reveals a sharp boundary -- the spectral edge -- between coherent optimization directions and stochastic noise, identified via the maximum consecutive singular value ratio σk/σk+1σ_k / σ_{k+1}. Across a 51M-parameter TinyStories model (4 seeds) and GPT-2 124M under distribution shift, the spectral edge exhibits a universal three-phase pattern (rise, plateau, collapse). The effective signal rank adapts to task complexity (k∗=2k^* = 2 at 51M, k∗=3k^* = 3 at 124M), and the directional coupling between spectral geometry and validation loss reverses with window size -- a lag flip reflecting the timescale of trajectory integration. Johnson--Lindenstrauss projection to d=10Wd = 10W dimensions (e.g., d=100d = 100 for W=10W = 10) preserves the spectral gap within 5.7%5.7\%, making the framework applicable to models of arbitrary scale. In companion work, the same spectral geometry provides early-warning signals of grokking -- predicting generalization 600600--1,7001{,}700 steps before it occurs across modular arithmetic, Dyck languages, and the SCAN benchmark.

Yongzhong Xu

Transformer Training DynamicsNeural Network Training DynamicsGradient Descent DynamicsSpectral Methods
AskPDF
Mar 10, 2026

Large Spikes in Stochastic Gradient Descent: A Large-Deviations View

Large loss spikes in stochastic gradient descent are studied through a rigorous large-deviations analysis for a shallow, fully connected network in the NTK scaling. In contrast to full-batch gradient descent, the catapult phase is shown to split into inflationary and deflationary regimes, determined by an explicit log-drift criterion. In both cases, large spikes are shown to be at least polynomially likely. In addition, these spikes are shown to be the dominant mechanism by which sharp minima are escaped and curvature is reduced, thereby favouring flatter solutions. Corresponding results are also obtained for certain ReLU networks, and implications for curriculum learning are derived.

Benjamin Gess, Daniel Heydecker

Technische Universität Berlin, Fakultät II 10623 Berlin, Germany · Max Planck Institute for Mathematics in the Sciences 04103 Leipzig, Germany · University of Oslo, Blindern 0316 Oslo, Norway

Large Deviation PrinciplesGradient Descent DynamicsStochastic Gradient Descent
AskPDF
Mar 6, 2026

Gradient Flow Polarizes Softmax Outputs towards Low-Entropy Solutions

Understanding the intricate non-convex training dynamics of softmax-based models is crucial for explaining the empirical success of transformers. In this article, we analyze the gradient flow dynamics of the value-softmax model, defined as L(Vσ(a)){L}(\mathbf{V} σ(\mathbf{a})), where V\mathbf{V} and a\mathbf{a} are a learnable value matrix and attention vector, respectively. As the matrix times softmax vector parameterization constitutes the core building block of self-attention, our analysis provides direct insight into transformer's training dynamics. We reveal that gradient flow on this structure inherently drives the optimization toward solutions characterized by low-entropy outputs. We demonstrate the universality of this polarizing effect across various objectives, including logistic and square loss. Furthermore, we discuss the practical implications of these theoretical results, offering a formal mechanism for empirical phenomena such as attention sinks and massive activations.

Aditya Varre, Mark Rofin, Nicolas Flammarion

Theory of Machine Learning Lab, EPFL.

Attention EntropyGradient Descent DynamicsSoftmax AttentionTransformer Attention
AskPDF
Mar 5, 2026

The Norm-Separation Delay Law of Grokking: A First-Principles Theory of Delayed Generalization

Grokking -- the sudden generalisation that appears long after a model has perfectly memorised its training data -- has been widely observed but lacks a quantitative theory explaining the length of the delay. We show that grokking is a norm-driven representational phase transition in regularised training dynamics, and establish the Norm-Separation Delay Law: Tgrok−Tmem=Θ(γeff−1log⁡(∥θmem∥2/∥θpost∥2))T_{\mathrm{grok}} - T_{\mathrm{mem}} = Θ(γ_{\mathrm{eff}}^{-1} \log(\|θ_{\mathrm{mem}}\|^2 / \|θ_{\mathrm{post}}\|^2)), where γeffγ_{\mathrm{eff}} is the optimiser's effective contraction rate (γeff=ηλγ_{\mathrm{eff}} = ηλ for SGD, γeff≥ηλγ_{\mathrm{eff}} \ge ηλ for AdamW). The upper bound follows from a discrete Lyapunov contraction argument; the matching lower bound from dynamical constraints of regularised first-order optimisation. Across 293 training runs spanning modular addition, modular multiplication, and sparse parity, we confirm three falsifiable predictions: inverse scaling with weight decay (R2=0.97R^2 = 0.97), inverse scaling with learning rate (R2=0.92R^2 = 0.92), and logarithmic dependence on the norm ratio (Pearson r=0.91r = 0.91). A fourth finding reveals that grokking requires an optimiser capable of decoupling memorisation from contraction: SGD fails entirely at the same hyperparameters where AdamW reliably groks. These results reframe grokking not as a mysterious optimisation artefact but as a predictable consequence of norm separation between competing interpolating representations. We further derive a practical three-input algorithm that predicts grokking delay at memorisation time with 34.6% mean absolute error (bootstrap 95% CI [30.0%, 39.4%], N=60N=60 seeds), enabling principled early stopping.

Truong Xuan Khanh, Truong Quynh Hoa, Luu Duc Trung +1

H&K Research Studio, Clevix LLC, Hanoi, Vietnam · Banking Academy of Vietnam, Hanoi, Vietnam

GrokkingNeural Network GeneralizationGradient Descent DynamicsNeural Network Memorization
AskPDF
Mar 5, 2026

Non-Euclidean Gradient Descent Operates at the Edge of Stability

The Edge of Stability (EoS) is a phenomenon where the sharpness (largest eigenvalue) of the Hessian approaches and then hovers near the stability threshold 2/η2/η during gradient descent (GD) with step size ηη. Despite (apparently) violating classical smoothness assumptions, EoS has been widely observed in deep learning, but its theoretical foundations remain incomplete. We provide an interpretation of EoS through the lens of Directional Smoothness [Mishkin et al., 2024]. This interpretation naturally extends to non-Euclidean norms, which we use to define generalized sharpness under an arbitrary norm. Our generalized sharpness measure includes previously studied vanilla GD and preconditioned GD as special cases, as well as methods for which EoS has not been studied, such as ℓ∞\ell_{\infty}-descent, Block CD, Spectral GD, and their normalized versions. Through experiments on neural networks, we show that non-Euclidean GD with our generalized sharpness also exhibits progressive sharpening followed by oscillations around or above the threshold 2/η2/η. Practically, our framework provides a geometry-aware spectral diagnostic that can be applied across a broad class of non-Euclidean gradient methods.

Rustem Islamov, Michael Crawshaw, Jeremy Cohen +1

University of Basel, Switzerland · George Mason University, USA · Flatiron Institute, USA

Edge of StabilityGradient Descent Dynamics
AskPDF
Mar 2, 2026

On the Rate of Convergence of GD in Non-linear Neural Networks: An Adversarial Robustness Perspective

We study the convergence dynamics of Gradient Descent (GD) in a minimal binary classification setting, consisting of a two-neuron ReLU network and two training instances. We prove that even under these strong simplifying assumptions, while GD successfully converges to an optimal robustness margin, effectively maximizing the distance between the decision boundary and the training points, this convergence occurs at a prohibitively slow rate, scaling strictly as Θ(1/ln⁡(t))Θ(1/\ln(t)). To the best of our knowledge, this establishes the first explicit lower bound on the convergence rate of the robustness margin in a non-linear model. Through empirical simulations, we further demonstrate that this inherent failure mode is pervasive, exhibiting the exact same tight convergence rate across multiple natural network initializations. Our theoretical guarantees are derived via a rigorous analysis of the GD trajectories across the distinct activation patterns of the model. Specifically, we develop tight control over the system's dynamics to bound the trajectory of the decision boundary, overcoming the primary technical challenge introduced by the non-linear nature of the architecture.

Guy Smorodinsky, Sveta Gimpleson, Itay Safran

Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Israel · Department of Computer Science, Reichman University, Israel

Neural Network Training DynamicsGradient Descent DynamicsOptimization Convergence AnalysisAdversarial Robustness
AskPDF
Mar 2, 2026

Never Saddle for Reparameterized Steepest Descent as Mirror Flow

How does the choice of optimization algorithm shape a model's ability to learn features? To address this question for steepest descent methods --including sign descent, which is closely related to Adam --we introduce steepest mirror flows as a unifying theoretical framework. This framework reveals how optimization geometry governs learning dynamics, implicit bias, and sparsity and it provides two explanations for why Adam and AdamW often outperform SGD in fine-tuning. Focusing on diagonal linear networks and deep diagonal linear reparameterizations (a simplified proxy for attention), we show that steeper descent facilitates both saddle-point escape and feature learning. In contrast, gradient descent requires unrealistically large learning rates to escape saddles, an uncommon regime in fine-tuning. Empirically, we confirm that saddle-point escape is a central challenge in fine-tuning. Furthermore, we demonstrate that decoupled weight decay, as in AdamW, stabilizes feature learning by enforcing novel balance equations. Together, these results highlight two mechanisms how steepest descent can aid modern optimization.

Tom Jacobs, Chao Zhou, Rebekka Burkholz

CISPA Helmholtz Center for Information Security, Saarbrücken, Germany

Neural Network Training DynamicsGradient Descent DynamicsAdaptive Gradient MethodsGradient DescentMirror DescentRepresentation Learning
AskPDF
Mar 2, 2026

Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning

We study the training dynamics of gradient descent in a softmax self-attention layer trained to perform linear regression and show that a simple first-order optimization algorithm can converge to the globally optimal self-attention parameters at a geometric rate. Our analysis proceeds in two steps. First, we show that in the infinite-data limit the regression problem solved by the self-attention layer is equivalent to a nonconvex matrix factorization problem. Second, we exploit this connection to design a novel "structure-aware" variant of gradient descent which efficiently optimizes the original finite-data regression objective. Our optimization algorithm features several innovations over standard gradient descent, including a preconditioner and regularizer which help avoid spurious stationary points, and a data-dependent spectral initialization of parameters which lie near the manifold of global minima with high probability.

Gautam Goel, Mahdi Soltanolkotabi, Peter Bartlett

Simons Institute, UC Berkeley · Department of Electrical and Computer Engineering, University of Southern California

Self-AttentionSoftmax AttentionGradient Descent DynamicsGradient Descent
AskPDF
Feb 24, 2026

On the Convergence of Stochastic Gradient Descent with Perturbed Forward-Backward Passes

We study stochastic gradient descent (SGD) for composite optimization problems with NN sequential operators subject to perturbations in both the forward and backward passes. Unlike classical analyses that treat gradient noise as additive and localized, perturbations to intermediate outputs and gradients cascade through the computational graph, compounding geometrically with the number of operators. We present the first comprehensive theoretical analysis of this setting. Specifically, we characterize how forward and backward perturbations propagate and amplify within a single gradient step, derive convergence guarantees for both general non-convex objectives and functions satisfying the Polyak--Łojasiewicz condition, and identify conditions under which perturbations do not deteriorate the asymptotic convergence order. As a byproduct, our analysis furnishes a theoretical explanation for the gradient spiking phenomenon widely observed in deep learning, precisely characterizing the conditions under which training recovers from spikes or diverges. Experiments on logistic regression with convex and non-convex regularization validate our theories, illustrating the predicted spike behavior and the asymmetric sensitivity to forward versus backward perturbations.

Boao Kong, Hengrui Zhang, Kun Yuan

Peking University. · Sichuan University.

Stochastic Optimization ConvergenceStochastic Gradient DescentGradient Descent DynamicsNonconvex Stochastic Optimization
AskPDF
Feb 18, 2026

Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking

Grokking -- the delayed transition from memorization to generalization in small algorithmic tasks -- remains poorly understood. We present a geometric analysis of optimization dynamics in transformers trained on modular arithmetic. PCA of attention weight trajectories reveals that training evolves predominantly within a low-dimensional execution subspace, with a single principal component capturing 68-83% of trajectory variance. To probe loss-landscape geometry, we measure commutator defects -- the non-commutativity of successive gradient steps -- and project them onto this learned subspace. We find that curvature grows sharply in directions orthogonal to the execution subspace while the trajectory remains largely confined to it. Importantly, curvature growth consistently precedes generalization across learning rates and hyperparameter regimes, with the lead time obeying a power law in the grokking timescale. Causal intervention experiments show that motion along the learned subspace is necessary for grokking, while artificially increasing curvature is insufficient. Together, these results support a geometric account in which grokking reflects escape from a metastable regime characterized by low-dimensional confinement and transverse curvature accumulation. All findings replicate across this learning-rate range, a qualitatively different slow regime (lr=5e-5, wd=0.1, 3 layers), and three random seeds, though alignment dynamics differ quantitatively between regimes. Causal intervention experiments establish that orthogonal gradient flow is necessary but not sufficient for grokking: suppressing it prevents generalization with a monotonic dose-response across four operations, while artificially boosting curvature defects has no effect.

Yongzhong Xu

GrokkingNeural Network Training DynamicsGradient Descent DynamicsLoss Landscape Geometry
AskPDF
Feb 16, 2026

On the Stability of Nonlinear Dynamics in GD and SGD: Beyond Quadratic Potentials

The dynamical stability of the iterates during training plays a key role in determining the minima obtained by optimization algorithms. For example, stable solutions of gradient descent (GD) correspond to flat minima, which have been associated with favorable features. While prior work often relies on linearization to determine stability, it remains unclear whether linearized dynamics faithfully capture the full nonlinear behavior. Recent work has shown that GD may stably oscillate near a linearly unstable minimum and still converge once the step size decays, indicating that linear analysis can be misleading. In this work, we explicitly study the effect of nonlinear terms. Specifically, we derive an exact criterion for stable oscillations of GD near minima in the multivariate setting. Our condition depends on high-order derivatives, generalizing existing results. Extending the analysis to stochastic gradient descent (SGD), we show that nonlinear dynamics can diverge in expectation even if a single batch is unstable. This implies that stability can be dictated by a single batch that oscillates unstably, rather than an average effect, as linear analysis suggests. Finally, we prove that if all batches are linearly stable, the nonlinear dynamics of SGD are stable in expectation.

Rotem Mulayoff, Sebastian U. Stich

CISPA Helmholtz Center for Information Security, Saarbrücken, Germany

Gradient Descent DynamicsStochastic Gradient DescentDynamical SystemsNonconvex Optimization
AskPDF
Feb 9, 2026

On the Dynamics of Delayed Generalization in Linear Models

Delayed generalization, where training accuracy becomes nearly perfect substantially before test accuracy, has often been studied in the context of the depth and compositional structure of deep networks. In this paper, we ask whether delayed generalization can occur in simpler settings. We show that it can arise in a linear model trained with gradient descent and an exponential-tail loss on linearly separable data. We characterize the finite-time hyperplane dynamics into three phases: a population phase, a support vector dominated unlearning phase and a generalization phase. The convergence of the residual beyond the maximum-margin direction is sufficient for the unlearning phase to begin. In the support vector dominated phase, we characterize the hyperplane dynamics. The time scale in these dynamics causes delayed generalization. We validate our theory through controlled experiments varying population and support-vector distributions, and further verify the predicted bias behavior in a deep model with a single output neuron on MNIST. Our results provide a unified characterization of the asymptotic and finite-time dynamics underlying delayed generalization in linear models.

Atreya Vedantam, Nataraj Das, Chandrashekar Lakshminarayanan

Indian Institute of Technology Madras, Chennai.

Linear ClassificationNeural Network GeneralizationGradient Descent Dynamics
AskPDF
Feb 4, 2026

Gradient Flow Through Diagram Expansions: Learning Regimes and Explicit Solutions

We develop a general mathematical framework to analyze scaling regimes and derive explicit analytic solutions for gradient flow (GF) in large learning problems. Our key innovation is a formal power series expansion of the loss evolution, with coefficients encoded by diagrams akin to Feynman diagrams. We show that this expansion has a well-defined large-size limit that can be used to reveal different learning phases and, in some cases, to obtain explicit solutions of the nonlinear GF. We focus on learning Canonical Polyadic (CP) decompositions of high-order tensors, and show that this model has several distinct extreme lazy and rich GF regimes such as free evolution, NTK and under- and over-parameterized mean-field. We show that these regimes depend on the parameter scaling, tensor order, and symmetry of the model in a specific and subtle way. Moreover, we propose a general approach to summing the formal loss expansion by reducing it to a PDE; in a wide range of scenarios, it turns out to be first-order and solvable by the method of characteristics. We observe a very good agreement of our theoretical predictions with experimental results.

Dmitry Yarotsky, Eugene Golikov, Yaroslav Gusev

Applied AI Institute, Moscow, Russia · Steklov Mathematical Institute of Russian Academy of Sciences, Moscow, Russia.

Tensor DecompositionMean-Field TheoryNeural Tangent KernelGradient Descent Dynamics
AskPDF
Feb 1, 2026

Not All Preferences Deserve Gradients: Understanding Gradient Utility in Offline Reasoning Alignment

Offline preference optimization aligns reasoning models from fixed chosen--rejected pairs, yet standard methods apply gradient updates from every pair regardless of its training value under the current policy. We argue that this uniform treatment is wasteful and potentially harmful. From the perspective of gradient utility, we show that a pair's contribution depends jointly on informativeness and stability. Pair utility drifts as the policy evolves, high-gradient samples can coincide with high-curvature regions, leading to noisy and destabilizing updates, and the most effective supervision comes from stable confident errors where the model is reliably wrong yet curvature remains low. These findings motivate SAGE (Stability-Aware Gradient Efficiency), which maintains difficulty-stratified candidate pools refreshed during training and selects pairs within each pool by a forward-pass signal-to-curvature score. Only pairs with high current utility receive gradient computation; the rest are excluded from backpropagation. On mathematical reasoning benchmarks across multiple model scales, SAGE outperforms full-data and size-matched baselines while producing substantially smoother optimization trajectories.

Hui Wu, Hengyi Cai, Jinman Zhao +6

Baidu Inc. · University of Toronto · School of Intelligence Science and Technology, Peking University +1

Pairwise Preference LearningTraining Data SelectionPreference OptimizationOffline Preference OptimizationGradient Descent DynamicsGradient Alignment
AskPDF
Jan 30, 2026

Spectral Gradient Descent Mitigates Anisotropy-Driven Misalignment: A Case Study in Phase Retrieval

Spectral gradient methods, such as the Muon optimizer, modify gradient updates by preserving directional information while discarding scale, and have shown strong empirical performance in deep learning. We investigate the mechanisms underlying these gains through a dynamical analysis of a nonlinear phase retrieval model with anisotropic Gaussian inputs, equivalent to training a two-layer neural network with the quadratic activation and fixed second-layer weights. Focusing on a spiked covariance setting where the dominant variance direction is orthogonal to the signal, we show that gradient descent (GD) suffers from a variance-induced misalignment: during the early escaping stage, the high-variance but uninformative spike direction is multiplicatively amplified, degrading alignment with the true signal under strong anisotropy. In contrast, spectral gradient descent (SpecGD) removes this spike amplification effect, leading to stable alignment and accelerated noise contraction. Numerical experiments confirm the theory and show that these phenomena persist under broader anisotropic covariances.

Guillaume Braun, Han Bao, Wei Huang +1

RIKEN AIP · The Institute of Statistical Mathematics · Tohoku University +2

Gradient DescentGradient Descent DynamicsGradient Alignment
AskPDF
Jan 29, 2026

Riemannian Lyapunov Framework: Optimization as Closed Loop Control on Riemannian Manifolds

We introduce the Riemannian Lyapunov Framework (RLF), a unified framework that represents first order optimization as a closed loop controlled dynamical system on a Riemannian parameter manifold. Within this framework, optimizers are specified by a quadruple (M,g,Φ,η)(M, g, Φ, η) consisting of a Riemannian metric, an internal state transition map, a direction field generator, and a lifting parameter. Classical and modern methods including SGD, Adam, AdamW, Lion and Shampoo are recovered as specific instantiations of this quadruple. The central structural object of the framework is a time indexed family of target submanifolds in the extended state space of parameters, velocities, and internal memory, which we call the Normally Attracting Invariant Manifold (NAIM) family. These submanifolds organize training dynamics into two timescales: a fast contraction of the velocity state onto the current target submanifold, followed by a slow descent along it. Using Riemannian backstepping, we construct a strict Lyapunov function on the extended state space and prove Uniform Ultimate Boundedness (UUB) of the optimization trajectory under a descent alignment assumption and bounded stochastic disturbances. Under the additional Polyak-Lojasiewicz (PL) condition, the Lyapunov function converges linearly to a bounded neighborhood of the optimum, defined as the drift of the target direction field between consecutive iterates after parallel transport. To validate the framework's generativity, we derive three concrete instantiations of the quadruple, each exercising a distinct design axis, and show on geometric diagnostics, image classification benchmarks, and language modeling that their empirical behavior matches the theoretical predictions. The framework casts optimizer design as controller synthesis, complementing heuristic perspectives by providing a constructive template with provable guarantees.

Yixuan Wang, Omkar Sudhir Patil, Warren E. Dixon

Department of Mechanical and Aerospace Engineering University of Florida

Stochastic Optimization ConvergenceGradient Descent DynamicsLyapunov-Based ControlRiemannian Optimization
AskPDF
Jan 29, 2026

Why β1=β2β_1 = β_2 Is Dynamically Special in Adam

Adam has been at the core of large-scale training for almost a decade, yet the role of its two momentum parameters remains poorly understood. Recent work shows that tying β1=β2β_{1}=β_{2} can preserve Adam's strong performance despite collapsing two memory scales into one, raising a basic question: what becomes dynamically special when the memories are tied? We identify a concrete mechanism. In the continuous-time limit, each normalized-update coordinate decomposes into a sign component, an explicit magnitude-lag term proportional to the difference between the two memory times, and additional transition, curvature, and nonlinear ratio terms. This lag channel vanishes exactly when β1=β2β_{1}=β_{2}, making the diagonal the unique regime in which this mismatch-induced response is structurally absent. A full-history discrete decomposition on real training gradients recovers this change in composition: tied updates are sign-dominated, whereas the lag term becomes substantial off the diagonal and leaves a comparatively small residual. Across six vision and language tasks, tied configurations also typically exhibit smoother update-norm trajectories. Overall, our results identify memory-scale mismatch as a concrete source of magnitude sensitivity in Adam and provide a mechanistic account of why tied momentum is dynamically distinctive.

Alberto Fernández-Hernández, Cristian Pérez-Corral, Jose I. Mestre +2

Universitat Politècnica de València Valencia, Spain · Universitat Jaume I Castelló de la Plana, Spain

Momentum MethodsAdaptive Gradient MethodsNeural Network Training DynamicsGradient Descent Dynamics
AskPDF
Jan 27, 2026

To Grok Grokking: Provable Grokking in Ridge Regression

We study grokking, the onset of generalization long after overfitting, in a classical ridge regression setting. We prove end-to-end grokking results for learning over-parameterized linear regression models using gradient descent with weight decay. Specifically, we prove that the following stages occur: (i) the model overfits the training data early during training; (ii) poor generalization persists long after overfitting has manifested; and (iii) the generalization error eventually becomes arbitrarily small. Moreover, we show, both theoretically and empirically, that grokking can be amplified or eliminated in a principled manner through proper hyperparameter tuning. To the best of our knowledge, these are the first rigorous quantitative bounds on the generalization delay (which we refer to as the "grokking time") in terms of training hyperparameters. Lastly, going beyond the linear setting, we empirically demonstrate that our quantitative bounds also capture the behavior of grokking on non-linear neural networks. Our results suggest that grokking is not an inherent failure mode of deep learning, but rather a consequence of specific training conditions, and thus does not require fundamental changes to the model architecture or learning algorithm to avoid.

Mingyue Xu, Gal Vardi, Itay Safran

Department of Computer Science, Purdue University, West Lafayette, IN, USA · Department of Computer Science and Applied Mathematics, Weizmann Institute of Science, Israel · Stein Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Israel

Neural Network GeneralizationRidge RegressionGrokkingGradient Descent DynamicsGeneralization Bounds
AskPDF
Dec 31, 2025

Gradient Descent as Implicit EM in Distance-Based Neural Models

Neural networks trained with standard objectives exhibit behaviors characteristic of probabilistic inference: soft clustering, prototype specialization, and Bayesian uncertainty tracking. These phenomena appear across architectures -- in attention mechanisms, classification heads, and energy-based models -- yet existing explanations often rely on loose analogies to mixture models or post-hoc architectural interpretation. We provide a direct explanation. For any objective with log-sum-exp structure over distances or energies, the gradient with respect to each distance is exactly the negative posterior responsibility of the corresponding component: ∂L/∂dj=−rj\partial L / \partial d_j = -r_j. The identity is algebraic, requiring only differentiability; it is a specialization of Fisher's identity, and its significance here is its address: standard neural objectives instantiate it without modification. The consequence is that gradient descent on such objectives performs generalized expectation-maximization implicitly, with responsibilities arising as gradients to be applied rather than auxiliary variables to be computed. This result unifies three regimes of learning: unsupervised mixture modeling, where responsibilities are fully latent; attention, where responsibilities are conditioned on queries; and cross-entropy classification, where supervision clamps responsibilities to targets. Our claims live at training time: the responsibility-weighted gradient dynamics recently documented in transformers follow from the objective's geometry. The in-context Bayesian computation that trained transformers perform at inference time is the endpoint of these dynamics, not their per-step content.

Alan Oursland

Expectation-MaximizationNeural Network Training DynamicsGradient Descent Dynamics
AskPDF
Oct 15, 2025

Exact Dynamics of Multi-class Stochastic Gradient Descent

We develop a framework for analyzing the learning dynamics of high-dimensional problems trained using one-pass stochastic gradient descent (SGD) with data from multiple anisotropic classes. Our main theorem provides exact expressions for quantities of interest, including the risk and the overlap with the true signal, in terms of a deterministic system of ODEs, valid in the high-dimensional limit. The theorem holds for a broad class of optimization problems and extends to settings where the number of classes grows with dimension. To illustrate its utility, we investigate in detail the effect of the data's anisotropic structure on the problems of binary logistic regression and least-squares (LS) loss. We study the LS in a linear multiclass setup and derive a learning-rate threshold that depends on the average eigenvalue of the covariance matrices. In the binary logistic regression, we study three cases: isotropic covariances, data covariance matrices with a large fraction of zero eigenvalues (denoted as the zero-one model), and covariance matrices with power-law spectra. We show that a structural phase transition occurs. In particular, for the zero-one model and the power-law model with sufficiently large power, SGD aligns more closely with values of the class mean that are projected onto the ``clean directions'' (i.e., directions of smaller variance). This is supported by analytical studies and numerical simulations, which show the exact asymptotic behavior of the loss in the high-dimensional limit. The effects of data anisotropy that we demonstrate are likely to hold beyond these examples and illustrate one application of the broader theorem that we prove.

Elizabeth Collins-Woodfin, Inbar Seroussi

Logistic RegressionHigh-Dimensional AsymptoticsGradient Descent DynamicsStochastic Gradient Descent
AskPDF
Sep 10, 2025

Modified Loss of Momentum Gradient Descent: Fine-Grained Analysis

We analyze gradient descent with Polyak (1964) heavy-ball momentum (HB) whose fixed momentum hyperparameter β∈(0,1)β\in (0, 1) provides exponential decay of memory. Building on Kovachki and Stuart (2021), we prove that on an exponentially attractive invariant manifold the algorithm is exactly plain gradient descent with a modified loss, provided that the step size hh is small enough. Although the modified loss does not admit a closed-form expression, we describe it up to O(hR)O(h^{\mathcal{R}})-errors for arbitrary finite order R\mathcal{R}, and prove global (finite "time" horizon) trajectory approximation bounds O(hR)O(h^{\mathcal{R}}). We then conduct a fine-grained analysis of the combinatorics underlying the memoryless approximations of HB, in particular, finding a rich family of polynomials in ββ hidden inside which include and lie coefficient-wise in between Eulerian and Narayana polynomials. We prove that these polynomials are hh-polynomials of certain graph-associahedra. As corollaries of the main results, we derive continuous modified equations of arbitrary finite approximation order (with rigorous bounds) and the principal flow that approximates the HB dynamics, generalizing Rosca et al. (2023). Approximation theorems cover both full-batch and mini-batch HB. The results shed new light on the main features of HB and outline a roadmap for similar analysis of other optimization algorithms.

Matias D. Cattaneo, Boris Shigida

Princeton University

Gradient Descent Dynamics
AskPDF
Jul 15, 2025

Gradient Descent on Logistic Regression: Do Large Step-Sizes Work with Data on the Sphere?

Gradient descent on logistic regression has many fascinating properties not fully captured by classical analyses. For linearly separable data, the iterates converge in direction to the maximum-margin separator regardless of how large the step size is. For non-separable data, however, cycles may occur even for step sizes below the local stability threshold of 2/λ∗2/λ^*, where λ∗λ^* is the largest eigenvalue of the Hessian at the solution. We ask whether requiring all examples to have equal norm is a sufficient condition for global convergence below this threshold. We prove that this is indeed true in one dimension. In higher dimensions, even equal-norm datasets can produce cycles, although global convergence holds for the population loss when the examples are uniformly distributed on the sphere. Finally, in one dimension, we relax the equal-norm assumption to an annulus condition that ensures global convergence below the stability threshold, suggesting a path towards analogous sufficient conditions in higher dimensions.

Si Yi Meng, Baptiste Goujaud, Antonio Orvieto +1

Inria Paris, Paris, France · Inria Saclay, Palaiseau, France · ELLIS Institute Tübingen; MPI for Intelligent Systems; Tübingen AI Center, Germany +1

Logistic RegressionGradient DescentGradient Descent Dynamics
AskPDF