BRAIn Lab
Why SGD is not Brownian Motion: A New Perspective on Stochastic Dynamics Igor Ignashin1 , Anna Radovskaya1,3 , Andrew Semenov1,3 , Egor Lopatin1 , Stanislav Olegovich Potapov1 , Aleksandr Kovalenko1,3 , Andrey Veprikov1 , Aleksandr Shestakov1 , Andrey Leonidov1,3 , Aleksandr Beznosikov1,2
arXiv:2605.22644v1 [cs.LG] 21 May 2026
1 Basic Research of Artificial Intelligence Laboratory (BRAIn Lab)
1
2 Innopolis University
3 P.N. Lebedev Physical Institute of the Russian Academy of Sciences
Stochastic Gradient Descent (SGD) is commonly modeled as a Langevin process, assuming that minibatch √ noise acts as Brownian motion. However, this approximation relies on a continuous-time limit and a η noise scaling that does not match the discrete SGD update at finite learning rate. In this work, we propose an alternative formulation of SGD as deterministic dynamics in a fluctuating loss landscape induced by minibatch sampling. Starting directly from the discrete update, we derive a master equation for the parameter distribution and obtain a discrete Fokker–Planck equation that differs from the standard Langevin form at order η 2 . Using this framework, we analyze SGD dynamics near critical points of the loss. We show that the behavior decomposes along the eigenbasis of the mean Hessian into qualitatively distinct regimes. In particular, nearly-flat directions do not admit a stationary distribution: the variance grows over time, corresponding to effective diffusion along valleys with a coefficient proportional to the learning rate. We provide empirical evidence supporting these predictions on neural network models in computer vision and natural language processing, observing a clear qualitative separation between confined and diffusive modes.
mains both practically indispensable and theoretically important. Despite its widespread use, a complete theoretical understanding of SGD dynamics is still lacking. A dominant line of work interprets SGD through the lens of statistical physics, modeling it as a noisy relaxation process similar to Langevin dynamics [Mandt et al., 2017; Simsekli et al., 2019; Xie et al., 2021]. In this view, stochastic gradients are treated as Gaussian perturbations around the gradient of the average loss, leading—under simplifying assumptions—to a stationary Gibbs distribution. While this framework provides useful intuition, it relies on assumptions such as Gaussian gradient noise and the existence of a stationary regime, which may not hold in practical deep learning settings. As discussed in Section 2, empirical and theoretical studies increasingly point to systematic deviations from this Brownian-motion-based picture. An important step toward understanding these deviations was made by Feng and Tu [2021a], who established
Introduction
Stochastic Gradient Descent (SGD) has been the primary optimization method for training neural networks since the early days of deep learning [Rumelhart et al., 1986]. Despite significant advances in model architectures, datasets, and computational scale, SGD remains a central component of modern machine learning pipelines [Bottou, 2010; Goodfellow, 2016]. It is widely used to train over-parameterized models such as Transformers [Vaswani et al., 2017] and large-scale language models [Team et al., 2025; Hernández-Cano et al., 2025], and serves as the foundation for many adaptive optimization methods, including Adam [Kingma, 2014] and more recent variants such as Muon [Jordan et al., 2024]. In fact, recent empirical studies [Srećković et al., 2025] indicate that for sufficiently small batch sizes, the performance of adaptive optimizers approaches that of vanilla SGD, reinforcing its role as a fundamental baseline. These observations highlight that, despite its simplicity, SGD re1
• Empirical validation. We validate these predictions on vision and language models by analyzing Hessian spectra and parameter variance along eigendirections. The observed dynamics are consistent with the theoretical predictions.
an inverse variance–flatness relation for SGD near solutions. Their results show that directions with smaller curvature exhibit smaller fluctuations, revealing a nontrivial connection between loss geometry and stochastic dynamics. However, this relation is derived under specific assumptions and does not provide a general dynamical framework for SGD. In this work, we revisit the foundations of SGD dynamics and show that, at finite learning rates, the standard Langevin approximation does not faithfully reproduce the discrete SGD dynamics. Rather than modeling SGD as Brownian motion in a fixed potential, we interpret it as deterministic dynamics in a fluctuating loss landscape, where randomness arises from minibatch sampling. Unlike physical systems, where noise is external and uncontrollable, SGD stochasticity is algorithmically generated and directly tied to the structure of the loss function. This distinction allows us to derive the stochastic dynamics from first principles, without relying on continuoustime approximations or phenomenological noise models. Importantly, our results do not contradict continuoustime diffusion approximations in the infinitesimal stepsize limit, but instead highlight discrepancies that arise at finite step size.
2
Related Work
Langevin-based views of SGD. A classical line of research models SGD as a discretization of Langevin dynamics, establishing connections to equilibrium statistical physics. In Mandt et al. [2017], SGD was interpreted as approximate posterior sampling near local minima, with minibatch noise acting as an effective temperature. Related work [Smith and Le, 2017] analyzed the interplay between batch size, learning rate, and noise scale, leading to the “noise-scale” rule. Within this framework, SGD has been associated with implicit regularization and a preference for flat minima [Hu et al., 2017; Chaudhari and Soatto, 2018; Zhu et al., 2018; Yang et al., 2023]. While this perspective has been influential, it relies on approximations such as Gaussian gradient noise and the existence of a stationary regime. These assumptions are not always justified in modern deep learning settings and can lead to discrepancies between the Langevin approxContributions. Our main contributions are as follows: imation and the actual discrete SGD dynamics. • Discrete Fokker–Planck framework. Starting from the exact SGD update, we derive a discrete Fokker–Planck equation governing the evolution of the parameter distribution. We show that standard Langevin-based approximations omit terms of order η 2 , including contributions of the same order as the retained diffusion term, which can lead to qualitatively incorrect predictions at finite learning rates.
Heavy-tailed and non-Gaussian noise. Several works show that minibatch-gradient noise can be non-Gaussian or heavy-tailed, leading to dynamics that differ from Brownian motion [Simsekli et al., 2019, 2020; Gurbuzbalaban et al., 2021; Nguyen et al., 2019]. These works typically replace Gaussian noise with alternative stochastic models, whereas we derive the dynamics directly from the discrete SGD update.
• Analysis near critical points. We analyze SGD dynamics in the vicinity of critical points of the loss function using a quadratic approximation. This yields explicit expressions for the variance of parameter trajectories in terms of Hessian statistics.
Anisotropy and structure of gradient noise. Another line of work emphasizes that minibatch noise is highly anisotropic. In Jastrze˛bski et al. [2018], the covariance of gradient noise was shown to align with the Hessian, suggesting that the interaction between noise and curvature plays a key role in determining SGD behavior. Further work [Zhu et al., 2018] demonstrated that anisotropic noise facilitates escape from sharp minima and biases optimization toward flatter regions. These findings indicate that SGD noise cannot be accurately described as isotropic Brownian forcing. Instead, its structure is closely tied to the geometry of the loss landscape.
• Unified picture of SGD dynamics. Our framework reveals distinct dynamical regimes in the Hessian eigenbasis. Directions with positive curvature exhibit bounded fluctuations consistent with the inverse variance–flatness relation [Feng and Tu, 2021a], while nearly-flat directions remain nonstationary and exhibit diffusive behavior.
2
Alternative dynamical formulations. Beyond Langevinbased models, several works have explored alternative descriptions of SGD dynamics. In Yaida [2020], finitewidth neural networks were shown to exhibit non-Gaussian fluctuations, providing a different perspective on stochastic training dynamics. In Feng and Tu [2021b], SGD was analyzed in terms of dynamical phases, depending on noise and data properties. Closely related to our work, Feng and Tu [2021a] established an inverse variance–flatness relation, linking parameter fluctuations to local curvature. These approaches emphasize that SGD operates in a non-equilibrium regime that is not fully captured by classical Langevin formulations. Positioning of this work. Our work differs from these lines of research in two key aspects. First, we derive the evolution of the parameter distribution directly from the discrete SGD update via a master equation, rather than starting from a continuous-time stochastic differential equation. Second, we interpret SGD noise as arising from fluctuations of the loss landscape itself, rather than as an external stochastic forcing term. This leads to a discrete Fokker–Planck description that captures finitestep effects and provides a unified explanation of the dynamical regimes observed in practice. Figure 1: Left: Brownian motion, modeled as a particle driven by random forces in a static potential. Right: 3 Setup and Problem Formulation deterministic motion in a fluctuating potential, which SGD with replacement. We consider stochastic gra- better reflects SGD with minibatch sampling. dient descent (SGD) with sampling with replacement, so that at each iteration the minibatch is drawn indepenHowever, this construction does not exactly match SGD. dently from the dataset. The parameter update is In the true update (1), stochasticity enters through the loss itself and therefore appears at order η, not wn+1 = wn − η∇Ln (wn ), (1) minibatch √ η. As a result, the Langevin approximation does not where wn denotes the parameters at step n, Ln is the reproduce the discrete SGD process exactly and omits minibatch loss, and η is the learning rate. terms that are relevant at finite step size (see Appendix A). A more faithful interpretation is that SGD describes deBrownian motion versus a fluctuating landscape. terministic motion in a fluctuating loss landscape: the A common intuition is to view SGD as Brownian mo- objective changes from step to step because the sampled tion in a fixed potential, that is, as a particle driven by minibatch changes. In this view, the randomness is not random forces while evolving in the average loss L̄. This an external forcing term but an intrinsic consequence of viewpoint leads to a Langevin-type update minibatch sampling. Figure 1 illustrates the distinction. √ wn+1 = wn − η∇L̄(wn ) − η ξn , (2) Probabilistic description. To study the collective dy√ namics of SGD, we introduce the parameter distribution where ξn is interpreted as a stochastic force. The η scaling is introduced to obtain a nontrivial continuous- pn (w) at iteration n. Its evolution is exactly described time limit and the corresponding Fokker–Planck equation.
3
by the master equation Z h i pn (v) EL δ w − v + η∇L(v) dv, (3) pn+1 (w) =
• Difference from Langevin scaling. Langevin √ approaches impose η scaling of the noise in order to preserve diffusion in continuous time. In discrete SGD, the expansion is organized in powers of η, and this changes which terms appear at a given order.
Rd
where δ(·) is the Dirac delta function and EL [·] denotes expectation over minibatch sampling. Equation (3) is exact for SGD with replacement because minibatches are sampled independently across iterations. Without replacement, consecutive updates become correlated, and the corresponding master equation is more involved.
4
• Role of cumulants. At order O(η 2 ), only the second moment of minibatch gradients enters. Higherorder cumulants appear only at higher orders. • Gaussian minibatch fluctuations do not imply Gaussian parameter dynamics. Even if minibatch fluctuations are approximately Gaussian, the induced parameter dynamics need not be Gaussian.
Main Theoretical Results and Analysis
Discrete Fokker–Planck equation. Proposition 1 (Discrete Fokker–Planck equation for SGD). Assume that minibatches are sampled independently with replacement, the minibatch loss L(w) is sufficiently smooth, and the required gradient moments are finite. Then the SGD update
Mismatch with Langevin dynamics. Although Eq. (4) is exact up to order η 2 , it differs from the standard Langevin formulation already at this order. We now make this discrepancy explicit in a one-dimensional setting. We consider the one-dimensional SGD update
wn+1 = wn − η∇Ln (wn )
θn+1 = θn − η∇Ln (θn ),
induces the following second-order expansion of the parameter density: pn+1 (w) = pn (w) + η
d X
(5)
with moment assumptions
∇k pn (w) ∇k L̄(w)
EL [∇Ln (θ)] = ∇L̄(θ), 2 EL [(∇Ln (θ))2 ] = D(θ) + ∇L̄(θ) .
k=1 d η2 X + ∇k ∇l EL [∇k L(w)∇l L(w)] pn (w) 2
(6)
Proposition 2 (Finite-step mismatch with Langevin dynamics). Under the update (5) and moment assump+ O(η ). tions (6), the discrete SGD dynamics differ from the (4) Langevin-based Fokker–Planck equation at order η 2 . In particular, the discrete expansion contains an additional Proposition 1 is proved in Appendix A, where we derive second-order term proportional to (∇L̄)2 , which is absent the result by expanding the exact master equation for in the standard Langevin formulation. the SGD transition kernel and truncating the resulting series at second order in η. Equation (4) is the central Starting from the SGD update (Eq. (5)) and using the theoretical object of this work. It describes the evolution moment assumptions (Eq. (6)), the probability density of the parameter distribution under SGD directly in dis- evolves as crete time, without imposing a continuous-time approx∞ X ηm m imation. pn+1 (θ) = ∇ EL [(∇L(θ))m ] pn (θ) , (7) m! k,l=1 3
m=0
Implications of Eq. (4). Several important concluwhich yields, to second order, sions follow immediately. p (θ) = p (θ) + η∇ (∇ L̄(θ)) p (θ) n+1 n n • Discrete nature of the dynamics. The equa tion is intrinsically discrete. In the strict limit 1 + η 2 ∇2 D(θ) + (∇L̄(θ))2 pn (θ) η → 0, the stochastic contribution disappears and 2 SGD reduces to gradient descent. + O(η 3 ). 4
(8)
By contrast, the Langevin approximation leads to 1 ∂t p(t, θ) = ∇ (∇L̄(θ)) p(t, θ) + η∇2 D(θ) p(t, θ) . 2 (9) Comparing Eq. (8) with Eq. (9), we observe that the discrete expansion contains the additional second-order contribution 1 2 2 η ∇ (∇L̄(θ))2 pn (θ) , 2
Langevin truncation. Indeed, if one instead uses the modified Langevin equation whose diffusion coefficient retains the additional (∇L̄)2 contribution from the discrete expansion, then the same calculation recovers the discrete SGD criterion ηΓ < λ(2 − ηλ) and the limiting variance Π(t → ∞) =
which is absent from the Langevin-based Fokker–Planck equation. Since this term is of the same formal order as the retained diffusion contribution, the Langevin approximation does not provide a consistent finite-step truncation of the discrete SGD dynamics. A modified Langevin model, q 2 dθ(t) = −∇L̄(θ(t))dt − η D(θ(t)) + ∇L̄(θ) dWt , (10) recovers the discrete result up to second order, but still differs at higher orders.
ηd . λ(2 − ηλ) − ηΓ
(16)
Thus, the modified Langevin description (10) agrees with discrete SGD at this order, whereas the standard Langevin approximation does not. Therefore, in the regime 2λ − ηλ2 < ηΓ < 2λ the standard Langevin approximation predicts a stationary state with finite variance, while the exact discrete SGD dynamics yields exponential growth of variance. This is a qualitative difference, not merely a quantitative correction. Note that the modified Langevin equation (10) predicts the same behavior as exact SGD.
Toy example. To illustrate the practical effect of this discrepancy, consider
Local analysis near a critical point. We now ap(11) ply Eq. (4) to SGD near a critical point of the average loss. Let v be such a point. In a neighborhood of v, we where G and H are independent Gaussian random vari- approximate the minibatch loss by a quadratic form: ables with moments d d X 1 X L(w) = L + G (w −v )+ Hij (wi −vi )(wj −vj )+. . . 2 0 i i i EL [G] = 0, EL [G ] = d, 2 i=1 i,j=1 (12) EL [H] = λ, EL [H 2 ] = λ2 + Γ. (17) Here Gi = ∇i Ln (v) are the stochastic gradient compoThen nents and Hij is the stochastic Hessian at v. ∇L̄(θ) = λθ, D(θ) = d + Γθ2 . (13) Averaging over minibatches gives For discrete SGD, the long-time variance is finite if d 1 X E [L(w)] = E [L ]+ EL [Hij ](wi −vi )(wj −vj )+. . . , 0 L L ηΓ < λ(2 − ηλ), 2 1 L(θ) = L0 + Gθ + Hθ2 , 2
i,j=1
and in that case ηd Πn→∞ = . λ(2 − ηλ) − ηΓ
since EL [Gi ] = 0. Thus, v is a critical point of the average loss, although it need not be a minimum of each (14) individual minibatch loss.
For the standard Langevin approximation, the corresponding condition is only ηΓ < 2λ with limiting vari- Mean and covariance of parameter trajectories. We characterize the parameter distribution by its mean, ance ηd measured relative to v, Π(t → ∞) = . (15) Z 2λ − ηΓ n µi = (wi − vi ) pn (w) dw, This discrepancy is not caused by the toy model itself, but by the second-order terms omitted in the standard 5
and covariance matrix Z Πnij = (wi − vi )(wj − vj ) pn (w) dw − µni µnj .
Main result near a critical point.
Proposition 3 (Variance dynamics in the mean-Hessian eigenbasis). Let v be a critical point of the averaged loss, ∇L̄(v) = 0, and assume the local quadratic approximaSubstituting the quadratic approximation (17) into Eq. (4) tion yields a closed evolution equation for Πnij (Appendix B). d d X 1 X Hij (wi −vi )(wj −vj ). L(w) = L0 + Gi (wi −vi )+ Mean-Hessian eigenbasis. The covariance dynamics 2 i,j=1 i=1 simplify substantially in the eigenbasis of the mean Hessian, Let λi be the eigenvalues of the mean Hessian EL [H]. d X Under the assumptions stated above, the covariance maEL [Hij ] = λk Oki Okj , trix is approximately diagonal in the mean-Hessian eigenk=1 basis, and its diagonal entries satisfy where λk are the eigenvalues and O is orthogonal. In n 1 − 1 − 2ηλi + η 2 EL [H̃ii2 ] this basis we define n Π̃ii = ηγ (λi + ϵ). (19) 2λi − ηEL [H̃ii2 ] d d X X H̃ij = Oik Ojl Hkl , Π̃nij = Oik Ojl Πnkl . (18) Here γ is independent of the learning rate, ϵ is a small k,l=1 k,l=1 regularizer, and n is the iteration index. Empirically, the rotated minibatch Hessian H̃ij is diago- The derivation of Proposition 3 is given in Appendix B. nally dominant, with fluctuations primarily affecting the We emphasize that this expression is approximate and relies on the assumptions stated above. In particular, diagonal entries due to SGD noise. the proportionality between gradient noise and curvaAssumptions. To obtain a tractable closed-form ex- ture is an empirical observation rather than a general pression, we adopt the following assumptions (see Ap- theoretical property. pendix B): Corollary 1 (Diffusive and confined eigendirections). • independence of fluctuations of the components of Under the assumptions of Proposition 3, SGD separates into distinct dynamical regimes in the mean-Hessian eigenH̃ij ; basis. 2 • approximate diagonality of EL [H̃ij ]; First, for nearly flat directions, λi ≈ 0, the covariance grows approximately linearly with time: • small step size, ηλmax < 1, where λmax is the largest eigenvalue of the mean Hessian; Π̃nii ∝ n. • empirical relation between gradient covariance and More explicitly, before saturation and for ηn|λi | ≪ 1, the Hessian: recursion reduces to Π̃nii ≈ γη 2 n (λi + ϵ). EL [Gi Gj ] ≈ γ EL [Hij ] + ϵ , These directions remain diffusive over the observed time horizon and correspond to broad valleys of the loss landThese assumptions reflect the empirical structure of mini- scape. batch noise and the standard finite-step stability condi- Second, for sufficiently sharp directions with positive curtion. vature and The structure of the covariance dynamics can now be 1 − 2ηλi + η 2 EL [H̃ii2 ] < 1, made explicit in the mean-Hessian eigenbasis. with constant γ and small ϵ.
the covariance approaches a finite limiting value. In the small-step regime this limit is 1 Π̃ii ≈ ηγ. 2 6
Finally, intermediate directions interpolate between these two limits and are described by the full expression in Proposition 3. Thus, within the local quadratic regime and over the observed time horizon, SGD near a critical point does not generally converge to a fully stationary distribution: sharp directions are confined, while nearlyflat directions exhibit persistent spreading.
5
Experiments
Small-scale validation. We first validate the assumptions and structural predictions of the theory in a smallscale NanoGPT/Shakespeare setting, where Hessian statistics and covariance dynamics can be measured directly. In particular, access to exact Hessians allows us to explicitly verify the assumptions underlying the theoretical derivation. Figure 2 compares the theoretical diagonal covariance profile from Eq. (19) with the empirical covariance in the mean-Hessian eigenbasis. Since the overall multiplicative constant γ is not fixed by the theory, the comparison is structural rather than absolute. We observe good agreement, including the predicted separation between saturating directions and directions that remain non-stationary over the observed time horizon. This supports both the assumptions and the resulting covariance dynamics derived in Proposition 3. Figure 2: Small-scale NanoGPT/Shakespeare validation. Left: theoretical prediction for the diagonal Empirical validation at scale. Having validated the covariance profile from Eq. (19). Right: empirical meaassumptions in a controlled small-scale setting, we now surement in the mean-Hessian eigenbasis. The comparitest the variance prediction of Proposition 3 on a larger son is structural, since the overall multiplicative constant NanoGPT model with 6.6M parameters trained on γ is not fixed by the theory. WikiText-2. Results. Figure 3 directly tests the stationary predicSetup. We construct a reference point w∗ by SGD tion of Proposition 3. The empirical plateau Π̂∞ i is aptraining followed by a short full-gradient refinement. We proximately independent of λi and scales linearly with approximate the top-20 sharp eigendirections vi of the the learning rate η, in agreement with the theoretical mean Hessian EL [H] using stochastic Lanczos, and launch prediction. Deviations from the Langevin approximation an ensemble of N = 50 independent SGD trajectories increase with η, consistent with the predicted finite-step from w∗ using sampling with replacement. corrections. The empirical covariance is measured as h i Π̂nii = Varj=1,...,N ⟨wn(j) − w∗ , vi ⟩ , (20) Discussion. Together, Figures 2 and 3 support the theoretical picture: the assumptions hold in a controlled which provides a direct estimate of the theoretical quan- setting, and the resulting predictions remain accurate at tity Π̃nii in Eq. (19). We evaluate three learning rates scale. In particular, sharp directions are confined with a η ∈ {0.001, 0.005, 0.010}. A single γ̂ is estimated at finite variance, while nearly flat directions remain diffu1 sive over the observed time horizon. Additional impleη = 0.001 using the asymptotic relation Π̃∞ ii ≈ 2 ηγ and then used without refitting for the other learning rates. mentation details are provided in Appendix C.
7
Figure 3: Discrete SGD vs. Langevin approximation: saturation level (NanoGPT 6.6M). Empirical plateau Π̂∞ i vs. Hessian eigenvalue λi for the top-20 sharp directions and three learning rates. Circles: empirical means; crosses: discrete prediction (Eq. 19); diamonds: Langevin prediction (Eq. (15)). The Langevin approximation increasingly underestimates the plateau at larger η, while the discrete prediction remains consistent with the data, highlighting the importance of finite-step corrections.
6
rate. In particular, it reveals correction terms that are not captured by standard Langevin-based models. Using this framework, we analyzed SGD dynamics near critical points of the loss landscape. We showed that the evolution of parameter variance is governed by the local Hessian structure and exhibits distinct regimes across eigendirections. Within the local quadratic regime, directions with large curvature exhibit approximately stationary variance, while nearly-flat directions remain nonstationary and diffuse over time. This leads to a decomposition of SGD dynamics into confined and diffusive modes. We provided empirical evidence for these predictions by comparing theoretical variance profiles with measurements from SGD trajectories in vision and language models. The observed behavior is consistent with the theoretical picture, with clear separation between saturating and non-saturating directions. Overall, our results suggest that SGD can be more accurately understood as deterministic dynamics in a fluctuating loss landscape rather than as Brownian motion in a fixed potential, especially at finite learning rates. This perspective provides a unified explanation of several empirical phenomena and highlights the importance of discrete-time effects in understanding optimization dynamics.
Limitations
Our analysis relies on several simplifying assumptions and practical constraints. First, the theory is based on a local quadratic approximation near a reference point, which may not capture global dynamics or strongly nonsmooth regions of the loss landscape. Second, we assume approximate independence and diagonality of minibatch Hessian fluctuations in the mean-Hessian eigenbasis; while supported empirically, this may not hold universally across architectures or training regimes. Third, the analysis is formulated for sampling with replacement, which differs from standard epoch-based training. On the empirical side, exact validation is limited to small models where dense Hessians can be computed, while larger-scale experiments rely on low-dimensional diagnostics and primarily probe sharp directions. Finally, the comparison is partly structural: the framework does not determine certain constants (e.g., γ), and reproducibility is limited by the experimental pipeline and uncontrolled sources of randomness. Extending the analysis beyond local regimes, relaxing structural assumptions, and improving large-scale validation remain important directions for future work.
7
Conclusion
In this work, we developed a discrete-time description References of SGD dynamics by deriving a Fokker–Planck equation directly from the update rule. This approach avoids Léon Bottou. Large-scale machine learning with stochascontinuous-time approximations and makes explicit the tic gradient descent. In Proceedings of COMProle of minibatch-induced fluctuations at finite learning STAT’2010: 19th International Conference on Com8
putational StatisticsParis France, August 22-27, 2010 Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cecista, Laker Newhouse, and Jeremy BernKeynote, Invited and Contributed Papers, pages 177– stein. Muon: An optimizer for hidden layers in neu186. Springer, 2010. ral networks, 2024. URL https://kellerjordan. Pratik Chaudhari and Stefano Soatto. Stochastic gradigithub.io/posts/muon/. ent descent performs variational inference, converges to limit cycles for deep networks. In International Andrej Karpathy. char-rnn. https://github.com/ karpathy/char-rnn, 2015. Conference on Learning Representations, 2018. URL https://openreview.net/forum?id=HyWrIgW0W. Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. Yu Feng and Yuhai Tu. The inverse variance–flatness relation in stochastic gradient descent is critical for finding flat minima. Proceedings of the National Academy Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Availof Sciences, 118(9):e2015617118, 2021a. able: http://yann.lecun.com/exdb/mnist, 2, 2010. Yu Feng and Yuhai Tu. Phases of learning dynamics in artificial neural networks: with or without mislabeled Qianxiao Li, Long Chen, and Cheng Tai. Stochastic modified equations and adaptive stochastic gradient data. arXiv preprint arXiv:2101.06509, 2021b. algorithms. ICML, 2017. Ian Goodfellow. Deep learning, 2016. Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as approximate bayesian Mert Gürbüzbalaban, Asuman Ozdaglar, and Pablo A. inference. Journal of Machine Learning Research, 18 Parrilo. Why random reshuffling beats stochastic gra(134):1–35, 2017. dient descent. Mathematical Programming, 2019. Mert Gurbuzbalaban, Umut Simsekli, and Lingjiong Thanh Huy Nguyen, Umut Simsekli, Mert Gurbuzbalaban, and Gaël Richard. First exit time analysis of Zhu. The heavy-tail phenomenon in sgd. In Internastochastic gradient descent under heavy-tailed graditional Conference on Machine Learning, pages 3964– ent noise. Advances in neural information processing 3975. PMLR, 2021. systems, 32, 2019. Jeff Z. HaoChen and Suvrit Sra. Random shuffling beats Alec Radford and Karthik Narasimhan. Imsgd after finite epochs. ICML, 2020. proving language understanding by generAlejandro Hernández-Cano, Alexander Hägele, ative pre-training. Technical report, OpeAllen Hao Huang, Angelika Romanou, AntoninAI, 2018. URL https://cdn.openai.com/ Joan Solergibert, Barna Pasztor, Bettina Messmer, research-covers/language-unsupervised/ Dhia Garbaya, Eduard Frank Ďurech, Ido Hakimi, language_understanding_paper.pdf. et al. Apertus: Democratizing open and compliant llms for global language environments. arXiv preprint David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by backarXiv:2509.14233, 2025. propagating errors. nature, 323(6088):533–536, 1986. Wenqing Hu, Chris Junchi Li, Lei Li, and Jian-Guo Liu. On the diffusion approximation of nonconvex stochas- Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in tic gradient descent. arXiv preprint arXiv:1705.07562, deep neural networks. In International Conference on 2017. Machine Learning, pages 5827–5837. PMLR, 2019. Stanislaw Jastrze˛bski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Umut Simsekli, Lingjiong Zhu, Yee Whye Teh, and Mert Gurbuzbalaban. Fractional underdamped langevin On the relation between the sharpest directions of dynamics: Retargeting sgd with momentum under dnn loss and the sgd step length. arXiv preprint heavy-tailed gradient noise. In International conferarXiv:1807.05031, 2018. ence on machine learning, pages 8970–8980. PMLR, 2020. 9
Samuel L Smith and Quoc V Le. A bayesian perspective on generalization and stochastic gradient descent. arXiv preprint arXiv:1710.06451, 2017. Teodora Srećković, Jonas Geiping, and Antonio Orvieto. Is your batch size the problem? revisiting the adam-sgd gap in language modeling. arXiv preprint arXiv:2506.12543, 2025. Kimi Team, Yifan Bai, Yiping Bao, Guanduo Chen, Jiahao Chen, Ningxin Chen, Ruijue Chen, Yanru Chen, Yuankun Chen, Yutian Chen, et al. Kimi k2: Open agentic intelligence. arXiv preprint arXiv:2507.20534, 2025. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. Wikipedia contributors. Caesar cipher — Wikipedia, the free encyclopedia. https://en.wikipedia.org/ wiki/Caesar_cipher, 2025. [Online; accessed 25-Sep2025]. Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum? id=wXgk_iCiYGo. Sho Yaida. Non-gaussian processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning, pages 165–192. PMLR, 2020. Ning Yang, Chao Tang, and Yuhai Tu. Stochastic gradient descent introduces an effective landscapedependent regularization favoring flat solutions. Phys. Rev. Lett., 130:237101, Jun 2023. doi: 10.1103/ PhysRevLett.130.237101. URL https://link.aps. org/doi/10.1103/PhysRevLett.130.237101. Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from minima and regularization effects. stat, 1050:21, 2018.
10
Appendix Supplementary Materials for Why SGD is not Brownian Motion: A New Perspective on Stochastic Dynamics
A
SGD Fokker-Planck equation and Langevin-like approach
Proof of Proposition 1. We prove the result by expanding the exact master equation (Eq. (3)) in powers of η and identifying the resulting terms up to second order. In this appendix we derive the exact recursion for the parameter distribution induced by one SGD step and then obtain its second-order expansion in the learning rate. We then compare this discrete expansion with the standard Langevin-inspired surrogate dynamics and show that the two descriptions differ already at order η 2 .
A.1
Exact recursion from the master equation
Our starting point is the master equation, which relates the parameter distribution at two successive iterations (Eq. 3): Z pn+1 (w) =
pn (v) EL [δ (w − v + η∇L(v))] dv.
We employ the representation of the delta function via its Fourier transform: Z dp ip⊤ a δ(a) = e , p ∈ Rd . (2π)d This yields:
Z pn+1 (w) =
(20)
(21)
dv dp
h i 1 ip⊤ (w−v) iηp⊤ ∇L(v) p (v)e E e . n L (2π)d
(22)
dv dp
1 ⊤ pn (v)eip (w−v) pi1 · · · pim EL [∇i1 L(v) · · · ∇im L(v)] . d (2π)
(23)
We now expand the exponential in η: pn+1 (w) =
∞ X (iη)m m=0
m!
Using:
Z
d X
Z
i1 ,...,im =1
dp ip⊤ (w−v) e pi1 · · · pim = (−i)m ∇i1 · · · ∇im δ(w − v), (2π)d
(24)
we obtain the recursion: pn+1 (w) =
∞ X ηm m=0
m!
d X
∇i1 · · · ∇im (EL [∇i1 L(w) · · · ∇im L(w)]pn (w)) .
(25)
i1 ,...,im =1
Important: Equation (25) is an exact discrete identity for SGD with independent minibatch sampling. No continuous-time approximation or truncation has been used.
A.2
Second-order Fokker–Planck approximation
To obtain a Fokker–Planck-type equation, we truncate the exact recursion (25) at second order in η. This requires: • sufficient smoothness of the loss, • existence of gradient moments, • small learning rate η. The resulting second-order truncation is given in the main text. 11
A.3
Langevin-inspired surrogate dynamics
We now compare the above exact discrete expansion with the standard Langevin-inspired approximation of SGD. Recall the SGD update: i wn+1 = wni − η∇i L(wn ). (26) The standard approach decomposes the loss into mean and fluctuations and models SGD as: d
i wn+1 = wni − η∇i L̄(wn ) −
√ X ij η C (wn )ξnj ,
(27)
j=1
where ξn is i.i.d. Gaussian noise:
j E[ξni ξm ] = δ ij δnm .
The covariance is matched via: X C ik (w)C kj (w) = ηEL [∇i L(v)∇j L(v)] − η∇i L̄(w)∇j L̄(w) =: Wij (w).
(28)
k
The corresponding master equation becomes: Z √ pn+1 (w) = pn (v) Eξ δ w − v + η∇L̄(v) + ηC(v)ξn dv.
(29)
Using Fourier representation: Z pn+1 (w) =
dv dp
1 ip⊤ (w−v+η∇L̄(v)) h i√ηp⊤ C(v)ξn i e Eξ e pn (v). (2π)d
(30)
1 ip⊤ (w−v)+iηp⊤ ∇L̄(v)−ηp⊤ W (v)p e pn (v). (2π)d
(31)
Averaging over Gaussian noise gives: Z pn+1 (w) =
dv dp
For small η this yields: pn+1 (w) = pn (w) + η
X k
A.4
1 X ∇k ∇l (Wkl (w)pn (w)) + . . . ∇k pn (w)∇k L̄(w) + η 2
(32)
k,l
Finite-step mismatch at order η 2
We now compare Eq. (32) with the exact expansion (25). At first sight the expressions are similar. However, writing: Wkl (w) = ηEL [∇k L(w)∇l L(w)] − η∇k L̄(w)∇l L̄(w), we see that the diffusion term in Eq. (32) is already of order η 2 . In contrast, the exact expansion (25) contains the second-order contribution: η2 ∇k ∇l (EL [∇k L(w)∇l L(w)]pn (w)) , 2 which can be decomposed as: EL [∇k L(w)∇l L(w)] = Dkl (w) + ∇k L̄(w)∇l L̄(w). Therefore, the exact discrete expansion contains the term: η2 ∇k ∇l ∇k L̄(w)∇l L̄(w)pn (w) , 2 which is absent from the Langevin-based recursion (32). This discrepancy arises because the Langevin approximation retains a diffusion term of total order η 2 while discarding other contributions of the same order. 12
A.5
Conclusion of the comparison
Thus, the difference between discrete SGD and its Langevin approximation is not merely interpretational. The mismatch appears already at order η 2 and can lead to qualitatively different dynamics. In contrast to the Langevin construction, the discrete formulation (25) keeps all contributions consistently organized in powers of η.
B
Motion near the Landscape critical point: complete calculations
In this appendix we derive the covariance dynamics used in the main text under a local quadratic approximation near a critical point of the averaged loss. The structure of the derivation is as follows. First, we introduce a quadratic approximation of the loss and derive closed equations for the mean and covariance of the parameter distribution. Next, we rotate these equations to the eigenbasis of the mean Hessian, which diagonalizes the deterministic part of the dynamics. Finally, we introduce simplifying assumptions that allow us to obtain explicit expressions for the variance and connect them to the main-text result. Our starting point is the Fokker–Planck equation, which is pn+1 (w) = pn (w) + η
d X
∇k pn (w) ∇k L̄(w)
k=1 d 1 X + η2 ∇k ∇l EL [∇k L(w) ∇l L(w)] pn (w) + . . . (21) 2 k,l=1
To study the local dynamics, we approximate the minibatch loss in a neighborhood of a critical point wc = v of the averaged loss, where ∇L̄(v) = 0, by a quadratic expansion: L(w) = L0 +
d X
d
Gi (wi − vi ) +
i=1
1 X Hij (wi − vi )(wj − vj ) + . . . 2
(22)
i,j=1
Here L0 , Gi , and Hij fluctuate across minibatches. The average loss is d
EL [L(w)] = EL [L0 ] +
1 X EL [Hij ](wi − vi )(wj − vj ) + . . . 2
(23)
i,j=1
so EL [Gi ] = 0. Thus, although individual minibatch losses may have nonzero gradients at v, the average loss has a stationary point at v. This distinction is important: stochasticity enters through fluctuations around this critical point. Under the quadratic approximation, the dynamics of the distribution can be fully characterized by its first two moments. Importantly, in this regime these moments form a closed system, which makes the analysis analytically tractable. We define the mean displacement from the critical point as Z n µi = (wi − vi ) pn (w) dw (24) and the covariance matrix Πnij =
Z
(wi − vi )(wj − vj ) pn (w) dw − µni µnj .
(25)
Substituting the quadratic approximation into the Fokker–Planck equation and integrating by parts, one obtains a closed system of recursion relations for the mean and covariance: µin+1 = µin − η
d X j=1
13
EL [Hij ] µjn ,
(26)
ij Πij n+1 = Πn − η
d X
d X kj 2 2 ij Πik E [H ] + E [H ] Π + η EL [Hik Hjl ] Πkl L L kj ik n n n + η Λn ,
k=1
where Λij n = EL
"
Gi +
(27)
k,l=1 d X
Hik µkn
k=1
Gj +
d X
Hjl µln
d X
# −
l=1
EL [Hik ] EL [Hjl ] µkn µln .
(28)
k,l=1
The covariance dynamics derived above contain many coupled terms between different coordinates. To simplify the analysis, we now pass to the eigenbasis of the mean Hessian. Specifically, we diagonalize the matrix EL [Hij ] as EL [Hij ] =
d X
λk Oki Okj ,
k=1
where Oij is an orthogonal matrix and λi are the eigenvalues of ⟨Hij ⟩L . This transformation diagonalizes the deterministic part of the dynamics and provides a natural coordinate system in which different curvature directions can be analyzed independently. In rotated basis G̃i =
d X k=1
Oik Gk ,
H̃ij =
d X
µ̃in =
Oik Ojl Hkl ,
d X k=1
k,l=1
Oik µkn ,
Π̃ij n =
d X
Oik Ojl Πkl n.
(29)
k,l=1
and the equations (26), (27) have the form (30)
µ̃in+1 = µ̃in − ηλi µ̃in . ij ij 2 Π̃ij n+1 = Π̃n − η(λi + λj )Π̃n + η
d X
2 ij EL [H̃ik H̃jl ] Π̃kl n + η Λ̃n .
(31)
k,l=1
To make further analytical progress, we introduce the following simplifying assumptions on the statistics of minibatch fluctuations in the rotated basis. First, we assume that different components of the gradient fluctuations are uncorrelated: EL [G̃i G̃j ] = di δij , EL [H̃ij ] = λi δij , (32) Second, we assume that the Hessian fluctuations are weakly correlated and can be expressed through their second moments: (33) EL [H̃ik H̃jl ] = λi λj δik δjl + Γik δij δkl + δil δjk − δij δik δjl . Here Γij characterizes the variance of fluctuations of the (i, j)-th component of the Hessian in the rotated basis. We first solve the recursion for the mean. Since the equation is linear and diagonal in the rotated basis, it can be solved explicitly, yielding µ̃in = (1 − ηλi )n µ̃i0 , where µ̃i0 is the initial displacement. Substituting this solution for the mean into the definition of Λ̃ij n , we obtain n n i j Λ̃ij n = (1 − ηλi ) (1 − ηλj ) (1 − δij )Γij µ̃0 µ̃0 + δij
d X
Γik (1 − ηλk )2n (µ̃k0 )2 + di δij .
(34)
k=1
We now turn to the covariance dynamics. We begin with the off-diagonal elements (i ̸= j), for which the recursion simplifies: 2 2 n n i j Π̃ij = (1 − ηλ )(1 − ηλ ) + η Γ Π̃ij (35) i j ij n + η (1 − ηλi ) (1 − ηλj ) Γij µ̃0 µ̃0 . n+1 This linear recursion can be solved explicitly, yielding n Π̃ij (1 − ηλi )(1 − ηλj ) + η 2 Γij − (1 − ηλi )n (1 − ηλj )n µ̃i0 µ̃j0 . n = 14
(36)
For the diagonal elements one has the equation 2 ii 2 Π̃ii n+1 = (1 − ηλi ) Π̃n + η
d X
2n k 2 Γik Π̃kk + η 2 di . n + (1 − ηλk ) (µ̃0 )
(37)
k=1
This equation is more complicated, however it can be solved in the matrix form. Let us introduce the diagonal matrix L̂ with elements Lij = (1 − ηλi )2 δij and the matrix Γ̂ with elements Γij . Then the solution is Π̃ii n =
d X (L̂ + η 2 Γ̂)n − L̂n ik (µ̃k0 )2 k=1
+η
(38)
# " d X 1 − (L̂ + η 2 Γ̂)n 2 k=1
1 − L̂ − η 2 Γ̂
dk . ik
Solutions (36),(38) have two contributions to the variance matrix Π̃ij n . The first one is related to the initial conditions, whereas the second one is diagonal (in the basis chosen). Now let us analyze this general result in context of numerical experiments made by us. Typically in our experiments the fluctuation matrix η 2 Γ are rather small (but not negligible) and the matrix Γij itself are nearly diagonal. We observe that one can safely neglect the non-diagonal elements of Γij . Also one can check, that the term containing initial condition are unimportant. Then the variance matrix Π̃nij is diagonal in the basis composed of ⟨Hij ⟩L eigenvectors. This is supported by our experiments (see the main text). In other words in this basis all the directions evolve independently and their behavior depends on the eigenvalues λi . For near-zero eigenvalues the (1 − ηλi )n quantity might be still near 1 up to step n, so one can expand the variance in 2 ηλi . After the simple algebra we obtain Πii n ≈ η ndk if ηn|λi | ≪ 1. These directions are non-stationary and the dynamics in these directions remain diffusive. Although the effective diffusion coefficient ηdk turns out to be very small, the large number of these nearly-flat directions leads to a noticeable divergence of individual trajectories. In the directions with large eigenvalues ηnλi ≫ 1 the system reaches the stationary state given by Πii n ≈
ηdi . 2λi − ηλ2i − ηΓii
(39)
The behavior in the other (intermediate) directions is in between these two limiting cases. We also observe empirically that, for directions with positive eigenvalues, di is approximately proportional to λi . In other words EL [Gi Gj ] ≈ γ EL [Hij ] + ϵij , where the matrix ϵij is small and necessary only for the regularization of negative eigenvalues. γ is a some constant, which is independent of the learning rate. We propose a possible mechanism of this phenomena below, but now we substitute this relation into the result for large λi directions Πnii ≈
ηγ ηγλi ≈ . 2 2 2λi − ηλi − ηΓii
(40)
We obtain the following picture. Near the critical point, there are many directions with small eigenvalues in which the system moves chaotically, resulting in divergence of trajectories. Whereas in the the rigid direction there is something like the evolution to the stationary state, but the width of this stationary distribution is almost the same for all (rigid) directions, i.e. independent of λi . For the intermediate directions the variance depends on the iteration number n. Under the assumptions stated above it is necessary to consider only diagonal elements of matrix L̂ + η 2 Γ̂, which are (41) (1 − ηλi )2 + η 2 Γii = 1 − 2ηλi + η 2 EL [H̃ii2 ]. The last equality follows from the definition of Γij . Putting everything together one can reproduce the main result of our paper from the equation (38) n 1 − 1 − 2ηλi + η 2 EL [H̃ii2 ] n Π̃ii = ηγ (λi + ϵ). (42) 2λi − ηEL [H̃ii2 ] 15
This result naturally explains the inverse Einstein relation discussed earlier in the literature. It is just the nonstationary effect. The effective fluctuations in the direction of small λi are small and the system does not evolve in this directions resulting in the small distribution width at a finite number of steps. At the end of this appendix let us make a few comments about the origin of ⟨Gi Gj ⟩L ∼ ⟨Hij ⟩L relation. One of the possible mechanisms is the following. Consider the minibatch loss near the critical point in the form ! NB d X 1 X (k) L(w) = Φ bi (wi − vi ) . (43) NB k=1
i=1
where NB is a minibatch size and ⃗b(k) is a random vector which is attributed to the given sample. Φ(x) - is some function. Expanding it around the critical point one obtains N
d
B X 1 ′ X (k) Φ (0) bi (wi − vi ) L(w) ≈ Φ(0) + NB
k=1 i=1
N
d
B X X 1 (k) (k) ′′ + Φ (0) bi bj (wi − vi )(wj − vj ) + . . . (44) 2NB
k=1 i,j=1
From this one can identify N
Gi =
B 1 ′ X (k) bi , Φ (0) NB
N
Hij =
k=1
so the relation
B 1 ′′ X (k) (k) bi bj . Φ (0) NB
(45)
k=1
′ (0) 2 Φ 1 2 Φ′ (0) Eb [bi bj ] = Eb [Hij ]. Eb [Gi Gj ] = NB NB Φ′′ (0)
(46)
holds. This is only an illustrative example, but it captures the basic mechanism. The fluctuations of both the gradient and the Hessian originate from the same source, so they may share the same information about the randomness of the system and be proportional to the same matrix.
C
Experimental setup
The code repository with implementation details and scripts for reproducing the experiments is available at: https://github.com/brain-lab-research/SGDiffusion. This section provides implementation details for the experiments used in the paper. The small-scale MNIST and Shakespeare experiments allow exact dense Hessian computation and are used to validate the structural assumptions behind Proposition 3: approximate diagonalization in the mean-Hessian eigenbasis, separation into sharp and flat regimes, and agreement between the predicted and measured covariance profiles. The large-scale NanoGPT experiment in Section C.4 complements these small-scale diagnostics by testing the sharp-direction saturation law on a 6.6M-parameter model, where only Hessian-vector products are feasible. Across the small-scale experiments, we use the same protocol: 1. train a model with SGD until the loss is substantially reduced; 2. refine the resulting checkpoint with full-gradient descent to obtain a reference point w∗ ; 3. compute Hessian information near w∗ and rotate SGD trajectories into the mean-Hessian eigenbasis; 4. compare the measured covariance with the theoretical prediction from Eq. (19).
16
0
50 100 150 200 250 300 350
50
100
100
10 2
| i|
0
150
10 4
200
10 6
250 300
10 8
350 0
50
100
150
200
i
250
300
350
400
variance matrix
Figure 4: MLP experiment. Left: eigenvalue spectrum of the mean Hessian. Right: variance matrix of SGD trajectories in the mean-Hessian eigenbasis.
C.1
Small-scale MNIST experiment
We first consider MNIST [LeCun et al., 2010] with a compact MLP containing 386 parameters. This model is small enough to compute dense minibatch Hessians exactly. After SGD training and GD refinement, we use the resulting point w∗ as the reference point for local analysis. We then compute minibatch Hessians near w∗ and form the mean Hessian EL [H]. Its eigenvalue spectrum is shown in Fig. 4 (left). The spectrum reveals that most eigenvalues are clustered near zero or slightly negative, with only a few directions exhibiting significant positive curvature (λmax ≈ 6). In this local region, the spectrum is dominated by wide, nearly flat directions, together with a small number of sharper directions. To connect the spectrum with SGD dynamics, we launch 1000 independent SGD trajectories from w∗ and project the displacements wn − w∗ onto the eigenvectors of EL [H]. The resulting covariance matrix in this basis is shown in Fig. 4 (right). Its near-diagonal structure supports the assumption that the mean-Hessian eigenbasis approximately decouples the dominant fluctuation modes. Figure 5 compares the predicted diagonal covariance profile from Eq. (19) with the measured profile in the same mean-Hessian eigenbasis. Since the overall factor γ is not predicted by the theory, the comparison is primarily structural. The predicted profile captures the qualitative separation between directions that saturate and directions that remain non-stationary over the observed horizon.
C.2
Small-scale Shakespeare experiment
To test whether the same phenomena appear beyond vision tasks, we repeat the analysis on a text modeling task. We train a simplified GPT [Radford and Narasimhan, 2018] model on the Shakespeare corpus [Karpathy, 2015]. Here we use a simplified symbolic prediction task in which each character is shifted by one position in the alphabet, similar to a Caesar cipher [Wikipedia contributors, 2025]. This keeps the problem simple enough for full Hessian analysis while preserving the sequential structure of the data. For the language model, we repeated the same set of experiments as conducted for the MLP, with learning rate η = 0.001. The spectral structure differs in scale — the maximum eigenvalue is substantially larger, λmax ≈ 217 (Fig. 6, left). Nevertheless, as in the vision case, most eigenvalues are near zero, indicating that the loss landscape is again dominated by wide flat valleys. The variance matrix (Fig. 6, right) is again close to diagonal in the Hessian eigenbasis. This shows that parameter fluctuations decouple along eigendirections, and the distinction between diffusive and rigid modes derived in 17
Figure 5: MLP experiment. Left: theoretical prediction for the diagonal elements of the variance matrix from Eq. (19). Right: experimental measurement of the same quantity. Note that the theory does not predict the value of the overall multiplicative constant γ; therefore, the magnitudes of the theoretical and experimental curves do not coincide. the previous section applies equally well here. Example trajectories confirming these behaviors are provided in Section C.3. Taken together, these results show that the structural assumptions of the theory—approximate diagonalization in the mean-Hessian eigenbasis and separation into sharp and flat regimes—also hold in a language-model setting.
C.3
SGD Trajectories in Different Eigendirections
To illustrate the distinct behaviors predicted by our theoretical framework, we examine individual SGD trajectories projected onto representative eigendirections of the mean Hessian. We report results for both the MLP (vision) experiment (Fig. 7) and the Shakespeare (language) model (Fig. 8). Diffusive directions (top panels). For both vision and language tasks, in eigendirections corresponding to near-zero or slightly negative eigenvalues, SGD trajectories display unconfined behavior. Variance grows steadily in time, consistent with diffusion along valley floors and in agreement with the Fokker–Planck analysis. Rigid directions (bottom panels). In directions with large positive curvature, trajectories remain localized around the minimum. The variance quickly saturates to a finite value controlled by the balance between curvature and minibatch noise. In sharp positive-curvature directions, the variance saturates to a finite value. Under the empirical relation di ∝ λi , this plateau becomes approximately independent of λi , consistent with the inverse variance–flatness relation discussed in the main text. Cross-task consistency. The similarity of behaviors between MLP and Shakespeare supports the view that this decomposition is not specific to a single architecture or data domain: the decomposition of SGD dynamics into diffusive and rigid modes is not tied to the specific architecture or data domain, but instead appears to be a general property of optimization in high-dimensional loss landscapes.
C.4
Large-Scale Validation: NanoGPT 6.6M on WikiText-2
This section provides implementation details for the large-scale validation experiment presented in Section 5. In particular, we describe the training protocol, spectral estimation, and trajectory measurements underlying Fig. 3.
18
0 102
200
400
600
800
200
10 1 10 4 | i|
0
400
10 7
600
10 10 10 13
800
10 16 0
200
400
i
600
800
1000
variance matrix
Figure 6: NanoGPT language model. Left: eigenvalue spectrum of the mean Hessian. Right: variance matrix of parameter trajectories in the mean-Hessian eigenbasis. Model and data. We use a NanoGPT transformer with embedding dimension 64, 2 attention heads, 4 layers, context length 256, and MLP ratio 4. The model is trained on WikiText-2 tokenized with the GPT-2 BPE vocabulary. We use the standard train/validation split. Training protocol. Starting from a random initialization, we train with SGD using learning rate ηtrain = 0.05, minibatch size B = 32, and sampling with replacement. Training proceeds until the validation loss plateaus: we stop when the mean loss over the most recent 200 steps improves by less than 10−3 relative to the preceding window. This occurs after approximately 5,150 SGD steps at validation loss 6.63. We save the resulting checkpoint wSGD . Starting from wSGD , we apply 100 steps of full-gradient descent with learning rate 0.01, accumulating gradients over the entire training set at each step. The result is our reference point w∗ . After 100 full-gradient refinement steps, the full-batch gradient norm at w∗ satisfies ∥∇L̄(w∗ )∥2 = 6.0 × 10−2 , with relative norm ∥∇L̄(w∗ )∥2 /∥w∗ ∥2 = 1.0 × 10−3 . As a measure of proximity to a critical point, we report ∥∇L̄(w∗ )∥2 /λmax = 1.1 × 10−3 ≪ 1, confirming that the gradient is small relative to the curvature scale and that the local quadratic approximation underlying Proposition3 is appropriate. Spectral analysis via stochastic Lanczos. Since exact dense Hessian computation is infeasible at this scale, we apply stochastic Lanczos with full reorthogonalization (every 10 iterations) to approximate the top-K = 20 eigenvectors of E[H] at w∗ . Each Hessian-vector product is averaged over 5 independent minibatches of size 4, and we run 200 Lanczos iterations in total. This yields eigenvalues λi ∈ [46.5, 55.8], all strictly positive. We note that w∗ is not a true minimum of the loss — as is typical for large neural networks, the Hessian at a post-training checkpoint has negative eigenvalues, indicating residual saddle structure. This does not affect the validity of the analysis in sharp directions, where Proposition3 applies, but prevents us from directly probing the diffusive regime at this scale. Local statistics at w∗ . We estimate the gradient noise covariance Kij = E[G̃i G̃j ] and the diagonal second moment E[H̃ii2 ] projected onto the top-20 sharp eigenvectors. Kij is computed by averaging outer products of per-minibatch gradients over Ngrad = 500 minibatches of size B = 32. E[H̃ii2 ] is estimated by averaging squared Hessian-vector products (vi⊤ Hn vi )2 over NH2 = 300 independent minibatches of size 4; the reduced minibatch size 19
Figure 7: MLP experiment: trajectories in distinct eigendirections of the mean Hessian. Top: flat directions (λ ≈ 0) show diffusive motion, with variance increasing over time. Bottom: curved directions (λ > 0) remain confined, with variance saturating at a finite value. is due to memory constraints, as the double-backward pass required for HVP computation increases peak GPU memory by approximately 2–3× relative to a standard forward pass. Estimation of γ. We estimate γ from the saturation level of the parameter variance ensemble at η = 0.001, 1 where finite-step corrections η E[H̃ii2 ] are negligible relative to 2λi (below 0.1%), so that Π̂∞ ii ≈ 2 ηγ holds to high accuracy. Specifically, 2 γ̂ = Π̂∞ , (47) η where Π̂∞ is the mean empirical plateau averaged over the top-20 sharp directions and the last 200 trajectory steps. This yields γ̂ = 1.81 × 10−4 . As an independent check, a weighted least-squares regression of di ≈ γλi over sharp directions gives γ̂WLS = 1.47 × 10−4 with coefficient of variation CV = 0.23, consistent with the saturation estimate within the expected uncertainty. SGD trajectory ensemble. We generate N = 50 independent SGD trajectories from w∗ , each using sampling with replacement and minibatch size 32. At each step n, we project the parameter displacement wn − w∗ onto all 20 sharp eigenvectors and record the scalar projections. We repeat this for three learning rates η ∈ {0.001, 0.005, 0.010}, running T = 500 steps for η ∈ {0.001, 0.010} and T = 1000 steps for η = 0.005. The empirical variance in direction i at step n is estimated as h i Π̂nii = Varj=1,...,N ⟨wn(j) − w∗ , vi ⟩ , (48)
20
Figure 8: Shakespeare model: trajectories in distinct eigendirections of the mean Hessian. Top: flat directions exhibit diffusive behavior. Bottom: curved directions show confined dynamics, consistent with the MLP case. where the variance is taken across the N trajectory realizations, providing a direct test of the predicted variance dynamics (see Fig. 9). Results. The main empirical result is shown in Fig. 3 in the main text. It demonstrates that the saturation level 1 ∞ Π̂∞ i is approximately independent of λi , in agreement with the theoretical prediction Π̃ii ≈ 2 ηγ. At η = 0.001, the saturation curve exhibits a clearly visible exponential approach with τ ≈ 9–11 steps, and the empirical trajectory closely tracks the theoretical prediction of eq.19. At η = 0.005 and η = 0.010, the saturation is nearly instantaneous (τ ≈ 2 and τ ≈ 1 step respectively). The simple approximation 12 ηγ̂ (Corollary 1) deviates from the empirical plateau by up to 25% at η = 0.010, while the full eq.19 — which retains the η E[H̃ii2 ] correction — remains within the empirical 90% bootstrap CIs at all three learning rates, with a median relative error below 10%. Figure 10 demonstrates that the mean saturation level scales linearly with η across one order of magnitude, as predicted by Corollary1. The theoretical line 21 γ̂η with γ̂ fixed from η = 0.001 achieves R2 = 0.90 against the three empirical measurements, with the largest relative error at η = 0.010 (25%) consistent with the growing importance of finite-step corrections at larger learning rates. Discrete SGD vs. Langevin approximation. We further compare the discrete prediction of eq.19 against the standard Langevin approximation ΠLang = ηdi /(2λi − ηΓii ) (eq. (15)), which differs from eq.19 only by the absence i 2 of the −ηλi term in the denominator. Using the same γ̂ and the same noise ansatz di = γ̂λi for both predictions, Fig. 3 shows that at η = 0.001 the two predictions agree within 3%, while at η = 0.010 the Langevin approximation systematically underpredicts the empirical plateau by ∼23% — an empirical confirmation at 6.6M-parameter scale of the finite-step mismatch predicted by Proposition 4.2.
21
Taken together, these implementation details support the empirical validation presented in the main text and clarify how the quantities entering Eq. (19) are estimated in practice at large scale.
D
Additional Experiments
In this section, we provide additional experimental results supporting the main claims of the paper. The experiments are designed to test four aspects of the theory: (i) whether continuous-time Gaussian approximations reproduce discrete SGD, (ii) whether the local covariance formula predicts the observed variance profile in the Hessian eigenbasis, (iii) whether flat and sharp directions exhibit distinct dynamical regimes, and (iv) whether the stochastic component of SGD can be causally controlled. Unless stated otherwise, all experiments use fixed random seeds, save the full configuration and raw intermediate artifacts, and generate figures directly from saved outputs. For empirical comparison, we construct discrete-time Gaussian surrogates driven by the full gradient and a minibatch-noise estimate. Importantly, the stochastic increment in these surrogates is scaled as O(η), matching the scale of the discrete SGD update, rather than adopting a literal continuous-time Brownian discretization with externally specified noise. As a result, these baselines should be interpreted as Gaussian approximations to SGD rather than exact realizations of a Langevin diffusion.
D.1
Quantitative Validation of the Discrete Theory
We next validate the local covariance prediction of Eq. (19). Starting from a trained reference point, we compute the full Hessian of the empirical loss and rotate trajectories into the eigenbasis of the mean Hessian. The theoretical prediction is computed from independently estimated gradient-noise statistics and Hessian eigenvalues; it is not obtained by fitting the observed trajectory variance. We perform this validation in two settings: the MLP-386 model on MNIST and the NanoGPT character-level language model on Shakespeare. In both cases, the reference point is obtained by training the model before the local Hessian and trajectory analysis are performed. This ensures that the covariance comparison is carried out in the same local regime as the one used in the theoretical derivation. For the MLP experiment, the predicted and measured variance profiles show strong structural agreement across Hessian eigendirections. In the full run, the log-correlation between predicted and measured variances is high, and held-out checks show that the agreement is not explained by a single post-hoc rescaling. Notably, the agreement is structural rather than merely scalar: the relative ordering of variances across eigendirections is preserved. We repeat the same validation for the NanoGPT model using the Shakespeare character-level task. The model has 960 parameters, allowing exact dense Hessian computation. The NanoGPT experiment follows the same protocol shape as the MLP experiment: the reference point is trained inside the experimental pipeline, the Hessian is computed at that point, and an ensemble of SGD trajectories is launched from the same initialization. The predicted and measured variance profiles show strong structural agreement, confirming that the covariance prediction is not specific to a single architecture or modality. We also estimate the proportionality constant γ across multiple sampler seeds and find it to be stable. Finally, we directly check the Hessian-noise structure in the mean-Hessian eigenbasis: the diagonal mass dominates the off-diagonal mass, supporting the approximate diagonality assumption used in the derivation.
D.2
Sampling Effects
Our theoretical master equation assumes sampling with replacement, ensuring independence of minibatches. In contrast, standard epoch-based training samples without replacement, introducing correlations between updates. Remark (sampling and stochastic modeling). The sampling rule is therefore part of the probabilistic model, not only an implementation detail. Langevin-type approximations usually rely on independent minibatch noise, which is naturally matched by with-replacement sampling [Li et al., 2017; Mandt et al., 2017]. In contrast,
22
reshuffling without replacement introduces correlations between updates and is known to modify the noise structure of SGD [Gürbüzbalaban et al., 2019; HaoChen and Sra, 2020]. To test the size of this effect in our setting, we start from the same MNIST/MLP-386 reference point and run trajectory ensembles under the two sampling schemes, keeping the optimizer, batch size, learning rate, and number of steps fixed. The trajectories are compared in the same projected coordinate system using ensemble variance, mean-path displacement, and the projected mean trajectory. We observe that sampling without replacement significantly suppresses the ensemble variance and introduces structured temporal correlations. This highlights that the sampling rule is part of the stochastic model and affects the validity of Langevin-type approximations.
E
Resources
All experiments were intentionally conducted on small-scale models and synthetic systems in order to make exact second-order diagnostics feasible and fully reproducible. For the MNIST experiments, we use an MLP with 386 trainable parameters. This allows exact dense Hessian computation by second-order automatic differentiation. A full Hessian has size 386 × 386, which is small enough to store and diagonalize directly on CPU. For the NanoGPT experiments, the model has approximately 960 trainable parameters. The corresponding dense Hessian has size 960 × 960. These experiments remain tractable, but dense Hessian computation is already the dominant cost. The additional reproducibility experiments in the exp6 suite were run on a local CPU machine with 12 CPU cores and 16GB RAM. CUDA was not used. The recorded full selected experiment pack completed in approximately one minute on this machine. Individual full runs ranged from below one second for small synthetic diagnostics to roughly 20 seconds for the largest ensemble simulations. The largest observed process RSS in the recorded runs was below 0.5GB, although this is only an end-of-run RSS snapshot rather than a true peak-memory measurement. Trajectory-based experiments can require substantially more storage when run at the larger settings described in the paper. For example, storing 1000 MLP trajectories of length 1111 with 386 parameters requires on the order of gigabytes. Similarly, storing many dense Hessians along a trajectory scales as O(T d2 ), where T is the number of Hessian evaluations and d is the number of parameters. All reported figures are generated automatically from saved intermediate artifacts, including raw trajectories, covariance statistics, Hessian spectra, and metric files. No manual figure editing or spreadsheet processing is required. These computational constraints are the reason the exact Hessian experiments are restricted to compact models. Extending the same diagnostics to larger modern architectures would require Hessian-vector products, Lanczos methods, or other approximate second-order techniques.
23
Figure 9: NanoGPT 6.6M on WikiText-2. Empirical variance Π̂nii (solid, with 90% bootstrap CI over N = 50 trajectories) vs. theoretical prediction of eq.19 (dashed) for three sharp eigendirections and three learning rates. The dotted line shows the Corollary1 approximation 12 ηγ̂. The theoretical time constant τi = (2ηλi )−1 is annotated per panel. A single γ̂ = 1.81 × 10−4 is estimated from η = 0.001 and used to predict all other curves. Shaded bands show 90% bootstrap confidence intervals obtained by resampling trajectories with replacement (150 resamples of N = 50 from 50), recomputing the variance estimate, and taking the 5th and 95th percentiles.
24
Figure 10: Mean saturation level Π̂∞ (averaged over 20 sharp directions) as a function of learning rate η. Left: linear axes. Right: log–log axes with the theoretical slope-1 line. The single γ̂ estimated from η = 0.001 predicts the other two points without refitting (R2 = 0.90).
Figure 11: Quantitative validation of the discrete theory on MLP-386/MNIST. Predicted and measured variances across eigendirections show strong structural agreement.
25
Figure 12: Quantitative validation of the discrete theory on NanoGPT/Shakespeare. The same covariance prediction remains structurally accurate for a small transformer language model.
Figure 13: Effect of sampling strategy on SGD dynamics. Both ensembles start from the same reference point and use the same optimizer hyperparameters; only the minibatch sampling rule is changed. Left: mean variance over the projected directions. Middle: distance of each mean path to the with-replacement SGD mean path. Right: first two coordinates of the projected mean trajectory. Sampling without replacement strongly suppresses the projected ensemble variance, while the mean trajectory remains close to the with-replacement trajectory.
26