Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold
Authors: Xinghan Li
Haodong Wen
Kaifeng Lyu
Venue: NeurIPS 2025 poster ()
Keywords: Adam
adaptive gradient methods
implicit bias
regularization
deep learning theory
Primary area: theory

Abstract: Despite the popularity of Adam optimizer in practice, most theoretical analyses study SGD as a proxy and little is known about how the solutions found by Adam differ. In this paper, we show that Adam reduces a specific form of sharpness measure shaped by its adaptive updates, leading to qualitatively different solutions from SGD.
    When the training loss is small, Adam wanders around the manifold of minimizers and takes semi-gradients to minimize this sharpness measure in an adaptive manner, a behavior we rigorously characterize via a continuous-time approximation using stochastic differential equations.
    We further illustrate how this behavior differs from that of SGD in a well-studied setting: when training overparameterized models with label noise, SGD has been shown to minimize the trace of the Hessian matrix, $\text{tr}(\textbf{H})$, whereas we prove that Adam minimizes $\text{tr}(\text{diag}(\textbf{H})^{1/2})$ instead. In solving sparse linear regression with diagonal linear networks, Adam provably achieves better sparsity and generalization than SGD due to this difference.
    Finally, we note that our proof framework applies not only to Adam but also to a broad class of adaptive gradient methods, including but not limited to RMSProp, Adam-mini, and Adalayer. This provides a unified perspective for analyzing how adaptive optimizers reduce sharpness and may offer insights for future optimizer design.

Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold
Xinghan Li' Haodong Wen"t Kaifeng Lyu'
Institute for Interdisciplinary Information Sciences
Tsinghua University
wxh-1i22,whd25) 0mails.............mmmmmmm..ll.3 wtt....ttsnsaa..aee edu. cn klyu@mail tww.gnua. edu
Abstract
Despite the popularity of the Adam optimizer in practice, most theoretical analyses study Stochastic Gradient Descent (SGD) as a proxy for Adam, and little is known about how the solutions found by Adam differ. In this paper, we show that Adam implicitly reduces a unique form of sharpness measure shaped by its adaptive updates, leading qualitatively different solutions from SGD. More specifically, when the training loss is small, Adam wanders around the manifold of minimizers and takes semi-g memi-gradients to minimize this sharpness measure in an adaptive manner, a behavior we rigorously characterize through a continuous-ting time approximation using stochastic differential equations. We further demonstrate how this behavior differs from that of SGD in well-studied setting: when training overparameterized models with label noise, SGD has been shown minimize the trace of the Hessian matrix, tr(H) H)), ), whereas we prove that Adam minimizes tr(Diag( H) 1/2) instead. In solving sparse linear regression with diagonal linear networks. this distinction enables Adam to achieve better sparsity and generalization than SGD. Finally, our analysis framework extends beyond Adam to a broad class of adaptive gradient methods, including RMSProp, Adam-mini, Adalayer and Shampoo, and provides a unified perspective on how these adaptive optimizers reduce sharpness, which we hope will offer insights for future optimizer design.
Introduction
Due to the non-convexity of the loss landscape, neural networks trained in different ways can perform very differently on the test set, even if they achieve the same training loss or accuracy (Zhang et al. 2017 Keskar al. 2017 Liu et al. 2023 Saunshi et al., 2024 To mathematically understand the generalization of neural networks, especially for over-parameterized models that admit many global minimizers, key step understand the implicit bias of optimization methods [Neyshabur et al.) 2014 Soudry et al.) 2018). That is, beyond just minimizing the training loss, what kinds of solutions are different optimizers implicitly biased towards?
Many theoretical works on implicit bias focus on (full-batch) gradient descent or its continuous variant, gradient flow. This includes the works the implicit bias towards max-margin classifiers (Soudry
et al., 2018, Nacson et al.] [019 Lyu and Li 2020 Ji and Telgarsky) 2020). implicit bias towards min-norm solutions [Lyu al. 2024), and equivalence to kernel methods [Jacot et al., 2018 Chizat jet al. 2019). However, these characterizations do not highlight the specific role of stochasticity in
SGID, although is more widely used in practice than gradient flow full-batch gradient descent.
"Equal contribution: alphabet ordering
Most work was done while Haodong was Xi' Ki'an Jiaotong University. "Correpponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Another line of works [Blanc et al. 2020 Damian et al., 2021 Li et al., 2021b] demonstrate that the gradient noise in SGD induces an additional form of implicit bias that reduces the sharpness of the solutions, a generalization measure that has been long observed to correlate with generaliza- tion [Hochreiter and Schmidhuber 1997. Keskar et al.) 2017 Jiang et al. 2020 Foret et al. 2021]. More specifically, these works focus on the dynamies of SGD when the training loss is already small and the iterates are close to a manifold of minimizers. Li et al. (022 introduced a general framework to analyze the dynamics of SGD near the minimizer manifold. showing that SGD will not stop at arbitrary global minimizers, but drift and diffuse around the manifold, driving the iterates towards flatter regions of the loss landscape.
This behavior mathematically characterized by a Stochastic Differential Equation (SDE), termed as slow SDE [Gu et al. 2023a]. This SDE accurately captures the projected dynamics of SGD near the minimizer manifold over timescale of (11 ") and reveals that SGD behaves like (emi-)gradient method on the manifold, taking semi-gradients to minimize specific sharpness measure determined by the Hessian and gradient noise. See|Section 3 for more details.
However, SGD is rarely used directly in modern deep learning. Instead, Adaptive Gradient Meth- ods (AGMs) have become the de facto standard for training neural networks. Among them, Adam Kingma and Ba 2014] innovatively combines the moving average of the first and sec- ond moments of gradients to determine an adaptive learning rate for each parameter. Across wide range of domains, Adam provides faster convergence and better stability than SGD [Ashish 2017
Dosovitskiy et al., 2020, |Schulman al. 2017 Zhang al., 2024c].
Despite the popularity of Adam, little is known about its implicit bias, especially how it is different from SGD in terms of reducing sharpness. In the literature, Ma et al. (2023) made attempts to generalize the slow SDE framework from SGD to Adam, but their analysis is specific to a two- dimensional loss function and involves a quasistatic approximation that lacks full mathematical rigor. Other works, such as Liu et al. [2023] and Gu et al. (2024), leverage insights from the slow SDE developed for SGD interpret empirical observations with Adam, but do not provide a theoretical analysis of Adam's 's own dynamics, A rigorous analysis of Adam's ''s implicit bias in terms of sharpness remains an open problem.
Our Contributions. In this paper, we show that Adam biases the iteration towards flatter regions in way that differs from SGD and implicitly reduces unique form of sharpness. We formally prove separations between SGD and Adam under concrete theoretical settings:
1. In [Section 4. we generalize the slow SDE for SGD Adam. Our slow SDE approximates the dynamics of Adam near the minimizer manifold, and reveals that Adam behaves like an adaptive gradient method that minimizes unique form of sharpness by taking semi-gradients on the manifold.
2. In Appendix we formally prove the generalization benefit of Adam when training overparam- eterized models with label noise. In this setting, we show that the implicit regularizer of Adam becomes tr tr(Diag( H 1/2), where H denotes the Hessian matrix. Compared to the implicit regularizer tr( H) of SGD in the same setting, this unique form of sharpness reduction can yield sparser solutions when the model parameterization aligns with the underlying problem structure. We empirically verify this theoretical prediction in the setting of solving sparse linear regression with diagonal lincar networks (Woodworth et al. 2020). where Adam recovers the sparse ground truth with less data. In Appendix [C] we further present a failure case of Adam's 's generalization: Adam does not outperform SGD in deep matrix factorization, likely because minimizing tr( tr(Diag( H) )/22 does not favor low-rank solutions.
3. On the technical side, our analysis framework extends beyond Adam to a broad class of adaptive gradient methods (AGMs), including Adam, RMSProp, Adam-mini, Adalayer, and Shampoo. We develop several techniques that may be of independent interest, including manifold projection operator tailored for AGMs, and high-probability convergence analysis of AGMs under
conditions that directly bounds L(Ox)) - L*
2 Related Work
Implicit Bias of SGD. A line of works studies the implicit bias of SGD when training overparame- terized models with label noise. Blanc et al. 2020) proved that with Eg loss and label noise, with

---

high probability, SGD will move away from points on the minimizer manifold I' if and only if tr (H) not locally minimal within I. HHoohhn et al. [202T] further extended this local characterization to global convergence result under the diagonal net setting [Woodworth et al. 2020]. However,
HaoChen et | (202T) relied on a manually designed, non-constant learning-rate schedule. [Damian et al.] [2021] overcame this issue by proving that the same implicit bias holds under constant learning
rates and for any smooth loss satisfying the Sardyka-Lojasiewiig condition. They also proved the
same implicit bias for SGD with momentum (SGDM).
Up to this point, all analyses were not able to track the optimization over On)=²) iterations, but this is necessary to capture the entire sharpness reduction process after converging to the minimizer manifold. [L et al. [2021b] were the first to tackle this problem, which provided an accurate SDE- characterization of the long-tew tengtterm behavior of SGD over O =3) time. Their SDE characterization also goes beyond the label noise case and is able to capture the implicit bias induced by general forms of gradient noise. However, the SDE derivation in Li et al. [2021b] is specific to plain SGD and cannot be directly extended to other optimizers. Gu et al. 2023b] later termed this form of SDE the slow SDE and derived it for local SGD [in et al. 2018 in a way that is potentially generalizable to broader class of optimizers. Wang et al. 2023] reinforced that momentum does not alter the implicit bias by showing that SGD and SGDM are characterized by the same slow SDE.
Another line of works characterizes the implicit bias of SGD through the lens of Gradient Regu- larization (GR). [L et al. [2018) first derived a stochastic modified equation of SGD, showing that introducing an additional gradient norm regularizer to the drift term of SDE helps approximate SGD more accurately. Later Barrett and Dherin [0020] also discovered the same regularizer that can be added to gradient flow to approximate GD in a higher order, and termed this behavior the Implicis Gradient Regularization (IGR). (mmihh et al. [022]] specifically studied and highlighted the IGR for SGD. Empirical studies [Geiping et al. [022 Novack et al. 2022] showed that GR is not necessarily implicit by demonstrating that explicitly regularizing gradient norm for SGD with larger or even full batch size can recover the generalization performance of small-batch SGD. GR was also found Minimization (SAM) [Foret et al. 2021] and inspired improvements upon SAM in terms of generalization performance [Zhao et al. 2022 The IGR method was later generalized to the analysis of as Adam and Adam by Cattaneo
Karakidda et al. [2023] to be connected with Sharpness-Awan
ct al.
2024] and [attanco and Shigida [2025). We further discuss our relationship to Cattanco et al.
[2024]] in the next paragraph.
Implicit Bias of Adam. On the theoretical side, the current literature still lacks a rigorous under- standing of the implicit bias of Adam, although it has been used more widely than SGD in practical deep learning, especially for large language models. However, many efforts have been made on this problem. Qian and Qian [2019] and Xie and Li] (2024) characterized the implicit biases of AdaGrad and Ad respectively. However, their methods do not extend to Adam. |Wang et al. [222]] showed that Adam's 's implicit regularizer is identical to that of SGD, while their result requires that the gradient coordinates have magnitude at most e, which is typically not feasible in practice. Zhang et al. [2024a]² 's analysis is also limited in scope, as it focuses on Adam's 's implicit bias on linearly separable data, a condition generally not met by real-world applications. A work using IGR to analyze the implicit bias of Adam, Cattaneo et al.) (2024), seems similar at first glance our study, but we are actually offering a different angle on the implicit bias of Adam. In particular, Cattaneo eet al. (2024) argued that full-batch Adam with constant learning rate approximately follows an ODE that anti-regularizes sharpness when 3, < 82. Our work analyzes the dynamics of Adam for O steps, longer horizon than Cattaneo et al. 2024 Our analysis shows that with gradient noise, Adam can be characterized by a slow SDE that regularizes sharpness in the long term, offering a complementary perspective.
Approximation of Stochastic First-Order Optimizers with SDE. Optimizers such as SGD and Adam take (1) steps to converge onto the manifold, and then move along the manifold for Onn steps, during which the optimizer will be dominated by slower implicit regularization dynamics, different from the convergence dynamics earlier. Before slow SDE was introduced by Li
et al. [2021b) the conventional SDE approximations [Li et al. 2018 202 Cattaneo et al. 2024
Malladi et al. 2024) track the iteration during the convergence phase, but they struggle to bound the approximation error if extended the manifold phase. Slow SDE tackles this problem by peeling the convergence dynamics off and approximating the iteration's 's projection on the manifold only. In this
way, slow SDE manages track iteration throughout the manifold phase for O(7 2) time. This idea will be made explicit in Section 3) In this work, to fill in the gaps and provide a theoretical analysis
3

---

that tracks iterations of Adam for sufficiently long time, we adopt the tool of slow SDE as in the aforementioned Li et al. 2021b] and Gu al. 2203a
Adaptive Gradient Methods. As a test-of-time optimizer that has revolutionized the field of deep learning [Kingma and Ba 2014], Adam innovatively combined the moving average of the first and second moments of gradients to determine an adaptive learning rate. Adam has also spawned a family of derivative optimizers such as AdamW [Loshchilov and Hutter, 2017]. AdaFactor (Shazeer and
Stern 2018 Adam-mini (Zhang et al. 2025]. Adalayer [Zhao et al. 2025 and AdaSGD Wang and
Wiens 2020 maintaining significant advantages over SGD in terms of empirical use. Under more
general framework of adaptive gradient methods. many optimizers also get huge success as adaptive gradient methods, such as RMSprop [Hinton et al. [012] and Adafactor [Shazeer and Stern [2118..
Convergence of Adam. There have been many previous works discussing the convergence bound of Adam. For example, seeddle,RReddi al. [20T8] and Dereich and Jentzen (2024) give convergence bounds
under the convexity condition, Zou et al. 2019 (Shi and Li] 2021] and Zhang et al. [2022] focus on
the cases where learning rates follow a 1/vl decay, and the bounds given by Zaheer et al. (2018).
/Zhang et a [2022] and /Wan et al. (2024b] do not decrease to 0 7 O. Also, most works
[Défossez et al. 2220 Guo et al. 2025 Liduka [022 Wang et al.] 2024a Zhang et al., 2022b] Hong
and Lin| 2023] only establish an upper bound on the average of gradient norms over the time of
iteration. In this work, we derive a high-probabilityy convergence analysis that directly bound the loss
term of the last step, L(Ox) f (1). Going beyond convex loss functions, we establish the
bound on μ-PL functions, and we focus on the constant learning rate schedule.
3 Preliminaries
Notations. Unless otherwise stated, for a square matrix M, we use diag (M) to denote the vector consisting of its diagonal entries. The notation Diag (with capital "D"" has two related usages: (1) for a vector v. Diag(v) denotes the diagonal matrix with v on its diagonal; and (2) for square matrix M, Diag(MM [M) denotes the diagonal matrix that keeps only the diagonal entries of M and zeros out the off-diagonal entries, i.e., Diag( M) == Diag(diag M)). For two vectors u, v with the same dimension d, u @ v denotes their element-wise product (W1V1,.... UddV)) For any exponent p > O, oP denotes element-wise exponentiation, i.c., o := (of...., 2)) and vuU stands for 001/2, We use IR to denote the set of d-dimensional vectors with non non-negative entries, and 8 to denote the set of d x d symmetric positive definite matrices. Given a pint O and the manimizer manifold T, we let (0st r) = min |0 41|2 be the &g distance between O and the manifold r.
Derivatives. For any scalar-valued function f : Rd + I. we write ( for its gradient at 0, and (I) for the projection of this gradient onto the tangent space of manifold r. For vector-valued function F Rd + Rdd we denote its Jacobian at 0 E I&d as OF(0) Rdxd, and its second-order derivative as 02 (0). which is a third-order tensor. Given a matrix M E dxdd we define 02 F() as the second-order directional derivative of F in the direction of M, 8F
2a [M Eveta
M) ei, where F; denotes the i-th coordinate of F, and e; the i-th 00
vector of the standard basis. When the context clear, for any scalar-valued function L Rd + R. we abbreviate 02 [M)[M] as [)
Loss Functions. Define E(0: 5) as the loss function for data sample & for model with parame- ters O. Define (())= := Eg-s (E( {)] as the training loss function, where S is the training dataset and ~ S means the data sample ç is drawn from S uniformly random. Let L = mingers L(0) be the minimum of training loss. Let Z(0) be the distribution of gradient noise Ve(0: {) (L))..
which is a random variable that depends on O. We define E() == Zzz ]] as the noise covariance matrix of gradients at 0.
Smoothness Assumptions. We make the following smoothness assumptions on the loss function and the gradient noise distribution.
Assumption 3.1. The loss function and the matrix square root of the noise covariance /1/2 are C -m-smooth on Rdd, i.e. all their partial derivatives up to order 5 exist and are continuous.
Assuming smoothness on the loss function is a common practice in optimization analysis. Here, we specifically assume the -o-smoothness, which we found to be a minimal smoothness requirement for our proof to hold for all C--smooth test functions in Theorem 4.1

---

Moreover, we assume that the smoothness constant of L and the gradient noise are globally bounded:
Assumption 3.2. L is p-smooth on R i.e. v 02 E R2, ([ - (2 112 - p| 02112 and C is bounded from below, i.e. & inf -o.
Assumption 3.3. The noisy gradients are Eg-bounded, i.e., there exists some constant R s.t. 40 R [[VE(0;))]l I R. almost surely for training data sample & Strain-
SGD and Adam. SGD is an iterative method that starts from an initial point Og and updates the
parameters as 0z+1 := 0₂ vl (0&) for all k 2 0, where 1j is the learning rate, (0) is the loss function for the data sample & sampled at step k. Adam (Kingma and Ba 2014] is a popular optimizer that updates the parameters as:
Mk+1 =- mimm + (1- (1
1k+ = 8202 + (1 - ()-) (ex(2)))
0.+1, == ON,5 (UK+1, +e
M
for all i [d].
Note that in practice, it is common to normalize m and Uk+1 by 1 - gftt and 1 1- s respectively before the division. However, this normalization quickly becomes neglectable when 7: is large, so we ignore it for simplicity.
SDE First-Order Approximation For SGD. A Stochastic Differential Equation (SDE) is an extension of an ordinary differential equation that incorporates random perturbations, and is widely used to model systems under the influence of noise. An SDE on Rd takes the form d0, = b(0,) ) (d,),,, where b: E + Rd is the drift vector field, IR + Rd is the diffusion matrix, and 1W)te is an mi-dimensional Wiener process. A line of works (Li et al. 2015 Jastrzębski et al. 2017 Li et al. 2017 (Smith et al., 2020 Li et al. 2019 2021a] used the following SDE to serve as first-order approximation of SGD, which we refer to as the conventional SDE:
d0,- = - (lt +
where the stochastic integral is taken in the Itô sense. For an introduction to Itô calculus, see Oksendal [2013). Later, Malladi et al. [2024] extended this type of SDE to Adam. Besides these conventional SDEs, below we introduce another type of SDE, slow SDE, that can more explicitly capture the implicit bias of SGD near a manifold of minimizers.
Manifold Assumption. Before going into the slow SDE, we introduce the manifold assumption. Previous studies [Garipov et al. 2018 Kuditipudi et al., 2019] have found that low-loss solutions are fact connected to each other, phenomenon known as mode connectivity. Wen et al. [0224] provided empirical evidence that the training dynamics of language model training usually happen
structure similar to a river valley, where many low-loss solutions lie the bottom of the valley.
Motivated by these observations, many previous works (Li al., 202 Fehrman et al. 2020 Lyu and Li 2020 Gu et al. [202aa] assumed that the minimizers of the training loss function are not isolated points but connected and form a manifold I::
Assumption 3.4. T is C -moothh, (d m) m)) dimenssioal compact submanifold of TRd, where any < T is a local minimizer of L. For all Ç E I, rank(V² = m. Additionally, there exists an open neighborhood T, denoted as U, such that I = arg mingeu L()).
With this assumption, if an optimization process converges and the learning rate 7j is sufficiently small, then the process will be trapped near some minimizer manifold which we denote by I..
Slow SDE. line of works [Baanc et al. 2020 Damian al. 2021 Li et al. 2021b) studied the dynamics of SGD near the manifold I' and showed that SGD has an implicit bias towards flatter minimizers on П. This effect cannot be directly seen from conventional SDEs, so et al. [002bb] derived a new type of SDE approximation, called slow SDE, that can explicitly capture this effect. See Appendix A for an illustration of the difference between conventional SDEs and slow SDEs. Here we introduce the slow SDE for SGD following the formulation in Gu et al. [2024] For ease of presentation, we define the following projection operators $ Pc for points and differential forms respectively. Consider the gradient flow dæ(t) - V ((x()) with z=0)
and fix some point Omall & T3 we define the gradient flow projection of any æ. (xx), as lim, ax(1) if the limit exists
and belongs to I, and mulll otherwise. It can be shown by simple calculus ei al. 2021b that a
5

---

equals the projection matrix onto the tangent space of I at Ç. We decompose the noise covariance Σ(() for Ç into two parts: the noise in the tangent space E == ( and the noise in the normal space E
Σ(() E))())
For any r, matrix A and vector b, we use ([AdW,+bii to denote ( + ((0())(+++(())))) [ dt by Itô calculus. Pc can be interpreted as projecting an infinitesimal step from Ç, so that Ç after taking the projected step does not leave the manifold r. Now we are ready to state the slow SDE for SGD.
which equals ( AdW, + (ab(C)b
Definition 3.1 (Slow SDE for SGD), Given > O and So E ET, define ((1) as the solution of the following SDE with initial condition ((0) Co:
dc(t) Pe (O)W:: V D (2(()))
(a) diffission
(b) drift
Here Eo(C) is defined as Eijz #o ,,, EE 0;0l) 1,og where {ui is an or- thonormal eigenbasis of 4 L(C with corresponding eigenvalues A1....,A4 Ad-
Interpretation of the Slow SDE for SGD: Semi meni-gradient Descent This SDE on the minimizer manifold r splits naturally into diffusion term Pc(E)//() dW;) injecting noise in the tangent (1=)//
space, and a drift term - 1 Pc (P))) [C) dt) that can be seen the negative semi-gradient of the following sharpness measure:
u(() = (2L(() (())))(((())).. (()
Here we use the word "semi-gradient" (Miih et al. 2015 Brandfonbrener and Bruna 2019 because it not exactly the gradient of μ(() but only the gradient with respect to the first argument of the inner product. More specifically, define y((1, (2) [(G,) Eo(C2) then the drift term is essentially - P(, (2) <1))<== after projecting onto the tangent space of I at Ç. In other words, SGD near manifold takes semi-gradients minimize the implicit regularizer (V² Eg (C)) but pretend (C( to be fixed, i.c. ignore the dependency Σ. (C) on Ç.
Example: Noisy Ellipse. We provide a toy example to illustrate the phenomenon described by the slow SDE for SGD: there are two parameters æ, y and an elliptical loss with label noise
[(x,y) = (x+z)²
[y=z)²
28 -1 The label noise ò is sampled uniformly from (-0.5, 0.5)
2a
every step. As depicted Fig. I, SGD moves towards flatter minimizers after reaching the manifold. The same phenomenon can be observed for Adam, but Adam converges to different minimizer that closer the axis (or, "sparser" in the parameter space). Understanding the difference between SGD and Adam is the main focus of this paper.
4 Theoretical Analysis of Adam
In this section, we generalize the slow SDE for SGD to general class of adaptive gradient methods (AGMs), including Adam. We first present our novel slow SDE for general class of AGMs, including Adam, and give an intuitive explanation for our results. Then, we discuss the difficulty of directly applying the slow SDE franework to Adam and other AGMs and how we resolve the problems.
General Class of Adaptive Gradient Methods.
We define general class of AGMs as follows:
m == B m. + (L)
Uk+ := 8r + (1 32) ( (V.(((.)VV
0 == 0 - -)S(#++1)
where d is the dimension of the parameter and the first order momentum m. and D is the dimension of the vector that encodes information related to second order momentum. For all optimizers but Shampoo in Table 1. we set d, while for Shampoo we have df + d2 if the matrix-like parameters have shape (d1, de) ). See Appendixx J for detailed explanation.

---

Contour
opom-SGD
x
Manifold
y
(a)
(b)
(c)
Figure I: (a): Coutour of the elliptical loss, from which we can see the two tips as the flattest minima. (b): SGD implicitly minimizes tr( (H) and converges to the flattest minima. (c): Adam reduces sharpness too but converges to a different and sparser minimizer.
The function S: RR20 - a p--smooth function for Ps 0, which maps non-negative
vector v _0 to a symmetric positive definite matrix S(v) E 84 addition, we require S to satisfy S(v) = for some Ro > O and any "2- We also require V Rdxdd _ RD to be linear function that satisfies Vggg ) E R 20 for all g Rd i.c., it always maps vector outer products to non-negative vectors.
number of currently used optimization algorithms, such RMSProp, Adam, Adam-mini, Adafac- tor', Adalayer, AdaSGD, and Shampoo?, all fit this framework. Note that we do not consider weight decays bias corrections in these optimizers. Some examples of V and S functions are listed in Table 1 including the AdamE-A optimizer that will be introduced in Appendix B as a tool to tune the implicit bias of Adam.
Prior to the results, we introduce two technical assumptions: S satisfies mild smoothness condition, and 1 31 is of constant order.
Assumption 4.1. The function S is C--soooth on R.D
Assumption 4.2. B1 0.9.
Remark 4.1. The threshold 0.9 in Assumption 4.2 can also be replaced by any constant below and the approximation rate in our result will remain unaffected. So we actually consider the regime where bj == 31 of constant order. For real-world Natural Language Processing (NLP) models, BERT (Devlin et al.) [2019], Transformer (Vaswani et al.) [017]] and GPT [Radford et al. [018]] all use 31 0.9. In computer vision (CV), pix2pix AIsola et al. [2017] uses 31 0.5, while U-Net [Ronneberger et al. y and ViT (Dosovisskiy et al. 2020 use 31 0.9. Thus, assuming 31 S 0.9 consistent with standard practice across multiple aspects.
4.1 Slow SDE Analysis for AGMs
Our SDE for AGMs characterizes the training dynamics near the manifold T. First rigorously define the preconditioned projection mapping 4 and the SDE projection formula as an extension to the & and Pt mentioned in Seccion 3] after which we present the SDE for AGMs derived.
Definition 4.1 (Preconditioner Flow Projection). Fix point 0mull & Given a Positive Semi-Definite --VV((x(I)) with xx x. We
matrix S. for E Rd, consider the preconditioner flow da(()
dt
denote the preconditioner flow projection of x as s i.e. $5(x) lim x(1) if the limit exists and belongs to F, and 4 s(x) = f otherwise.
We ignore update clipping, i.e. adopt the Algorithm 2 in Shazeer and Stern 2018
In practice, the Shampoo optimizer often equipped with the exponential moving average (EMA) on the calculation of condoiiiiioner [Morwani et al. 2024). Here we adopt this practical version of Shampoo instead of the original one [Gupta al. 2018)

---

Definition 4.2. For any Ç I and any differential form AdW, + bdt in Itô calculus, where A E Rdxd and R we use (c.s(AdW, + bdt) a shorthand for the differential form (C) (Obs(C)b + 100 J] dt.
Definition 4.3 (Slow SDE for AGMs). Given learning rate and the initial state Co E r, vo B q2
let c: 1-8a be a constant and denote St == S(v(t)), we define Ç(t) as the solution of the following
SDE with initial condition v())) = (Co. vo):
- ((() Σ ((()) - (((()))((())
dc(t) = P(()).())
(2)
diffissio
drift
du(() = c(V(Z(() - v)) dt.
preconditioner drife
Eo((:5)= SE(()S
En
(c: (())())((()))0
Note that the drift term in (t) can be interpreted an adaptive semi- remi-gradient descent process, in that this term drives the dynamics towards optimizing an adaptive loss function
u(C,v) = ((), ((())))(())
as if E. (S(t); S(t) has no dependence on Ç; also this gradient flow is preconditioned by positive definite matrix (t). Recall that the drift term in the slow SDE for SGD can be seen as a semi- gradient descent. In the AGM framework, it takes (1) time for the preconditioner S(l) to make a significant (i.e. ((1)) change, which coincides with the moving speed the slow SDE of Ç. Therefore, compared to that of SGD, our SDE includes preconditioner and injects adaptiveness accordingly in the semi- cond-gradient descent process.
formula that tracks the motion of the
We prove that ((t) always stays on the manifold r. And next, we present our main theorem showing that the above SDE in Equation (2) tracks the trajectory of Adam in a weak approximation sense.
Theorem 4.1. Let T and suppose Assumption 3./ 4.2 hold. There exist constant Ca. C >> such that the following statement holds for all sufficiently small learning rates Define step Kg == [olog(1//)) (oo)] and K == [Tn=² Run an AGM for Ko + K iterations and obtain {(0g, 2y 2b) 1
Ko+K [k=0 where the initial parameter Og is an arbitrary point that satisfies dist I' S cD. Let ((t)=
(C(1), v(()) be the solution to 2 on 1 E [0, T] with initial condition
(())= Xo T D
Xo: =
Vka).
For any is E {0. 10. K}. let Xx = ( (x) ()) (0Kg UKotk be the AGM state at step Ko + projected onto T. Then for any C functiin g(0),
B (xx) X E X (ka)))
= =
max OSKSK
where (()) hides logarithmic factors and constants independent of (but possibly depending on g)-
Theorem 4.1 shows that with small once AGM approaches the minimizer manifold, its long-
horizon behavior within (n) steps is captured by the SDE in Equation (2) A more explicit
version of Theorem 4. can be found in Appendix D
4.2 Interpretation of The Slow SDEs for AGMs
Adaptive Projection Operator. Equation (1) employs a fixed projection operator Pş to constrain the SDE to the manifold. As a comparison, the slow SDE for AGM uses an adaptive projection Pc.s(e) that depends on the current preconditioner S(v(1)). In other words, SGD's projection is state-independent, but AGM's projection is state-dependent. This adaptive projection alters the way the stochastic trajectory evolves on the manifold, giving rise to a different implicit bias AGMs versus SGD.
8

---

Table 1: Examples of optimizers the AGM Framework. See Appendix regularizers under label noise.
for derivations of their implicit
Implicit Regularizer (Label Noise,
Optimizer
Functions V/8
Remarks
V(M) 1a S: S(w) Ia
((H)) Blunc et
2020
V(M) diag(M)
Adam
S: S(v) Diag(
tr (Fr)
RMSProp
S: S(x) Ding 1/l//////////////////// <))
V(M) diag(M))
tr (Ern)(Er)
Adam with 03 0.
V(M): T E,en, M5s for E Bu
E.EMM v [B:]. tr (HF,)
Parameters partitioned into blocks (B1, B2. Bx ).
S(e) Diag(1/(v'
V(M), TEAT Eiccy Mas foralli La
Parameters are grouped by layess (Li, L2,
Adalayer
8: S(w) Ding((//(( (c)
EXEMM v /1.11- - (Hr,)
Lx }-
AdamE-A
V(M) diag(M)
u (Ding(H))*>))
See Section B. for discussion.
S(v) Diag(1/(e") +c)
V: V(M) (Vz(M), Vn(M))
[V.(M)),, = =Ex M.R.kki
For single matrix parameter. See Appendix J for discussion.
Shampoo
[V(MM),,,, = Ee Masua
explicit form
S:V V=)=((V=+=+]))) (VV
Effect of the Preconditioner on the Gradient Noise Covariance. Near the manifold, the gradient of loss vanishes (VL - 0), SGD's wandering around becomes noise-driven. For AGMs, the situation is more subtle.
First, one can show that the momentum term does not affect the implicit bias, consistent with prior theory [Wang et al., 2023]. The reason why 31 does not affect the implicit bias is that, after the iteration approaches the manifold, the difference between the current gradient gr and momentum M, becomes negligible in expectation.
Second, the AGM trajectory is influenced by its preconditioner. Concretely, the gradient- noise covariance matrix Σ is filtered through the preconditioner S(l) into (())E(() and then contributes to the SDE Over a long time horizon, this modified noise term alters the deterministic drift direction, further distinguishing AGM's dynamics from those of vanilla SGD.
4.3 Technical Difficulties and Proof Insights
4.3.1 Convergence Guarantee AGMs
The core of our study is to consider the behavior of Adam's implicit bias around the minimizer manifold. However, to make our study self-contained, we first need to show that Adam can actually converge to the neighborhood of the minimizer manifold under setting. This is non-trivial since Adam cannot provably converge to the minimizer manifold without any constraints. In fact, the convergence issue of Adam has been debated from its birth. Reddi et al. (0188 show that Adam does not converge to the optimal solution even in some simple convex settings. Recent work (Dereich and Jentzen 2024 gives Adam's ODE and shows that this ODE does not necessarily converge to the absorbing point of the gradient flow.
The magnitude of 1 .32 has been found to be important factor influencing whether Adam
converges. When - 3g is too large such that 3g < 88, Adam may not converge at all [Reddi et al. 2018 Zhang et al. 2024b). In our analysis, assume that Bg O i.c., 33 goes to zero as with a rate of r configuration we call the 2-scheme.
As stated in Sectiin 2] there have been works proving the convergence of Adam using various kinds of assumptions and giving different forms of bounds. However, previous results cannot be simply applied in our analysis. Specifically, our analysis near manifold requires high-probability bound on the optimality gap directly, and we have the 2-scheme assumption. Directly applying previous results in our setting yields bounds that are either loose or hold only in expectation, failing to meet

---

the requirements of our analysis. Moreover, we expect a bound that holds for all AGMs, instead of only Adam.
Therefore, we present the following statement of AGMs' convergence as a preparation for our subsequent study into their behavior the manifold.
Theorem 4.2 (Convergence Bound of AGMs, Stated Informally). Let Assumptions 3.2 3.3 and 4.2
hold, 32 ((T²), and L satisfy the Polyak- Lojasiewicz condition. With small learning rate ql. holds with high probability for some K =(1, log ()) that [(Ok) - O(1). See Theorem D.2
for a formal statement.
4.3.2 Key Insights in the Derivation of Slow SDEs for AGMs
After the AGMs reach the neighborhood of the minimizer manifold, we can derive an analysis similar the one in the local SGD paper [Gu al.) [023a]]. Specifically, we use SDEs to upproximate the AGMs after they reach the manifold neighborhood. However, unlike the usual SDE approximation, the SDEs we use here can track the AGMs for a much longer period of time, up to O(n rather than the O11 )) which is more common in the previous papers. This type of SDE is termed "slow SDE by Gu et al. 2023a
There are two obstacles preventing us from directly applying the analysis of slow SDEs from SGDs to AGMs. First, the obtaining of slow SDEs requires an accurate calculation of the variation of the first-order and second-order moments of the parameters over a relatively large number of steps (a "giant step" in the notation of Gu et al. [2023a]), and in the case of SGD, due to the nature of its rotational equivariance, we can always consider its Hessian matrix as a diagonal array, as well as its corresponding minimizer manifold as a space extended by some full-space standard bases, which greatly simplifies the computation. However, it is not the case for AGMs. Due to the effect of preconditioners S(xz), the rotation equivariance is not satisfied here.
To resolve this, generalize the gradient flow projection in Gu al. 2023a Li al. 202 into a varying preconditioner flow projection. Based on this definition, reparameterizing to the original space lets us reuse the simple formulas employed previously (Gu et al. 2023a Li et 2021b
The second reason is that when 32 is far from 1, the preconditioner changes too quickly, making the evolution of the moments hard characterize. Conversely, when B3 is extremely close to the preconditioner changes so little as to be impractical. Accordingly, we focus on the 2-scheme: 1 - B2 O(1²²) The key point is that this regime does not make the preconditioner's evolution negligible; rather, its slow but nontrivial drift shapes the SDE and can be tracked analytically.
Conclusions
In this work, have shown that Adam implicitly minimizes distinctive sharpness measure tr(Diag(H) H/2)) and that this bias leads to different solutions and generalization behavior compared to SGD. Our slow SDE framework not only rigorously characterizes Adam's adaptive semi-gradies drift near the minimizer manifold, but also recovers concrete separations in sparse linear regression and deep matrix factorization settings.
Despite these advances, several important avenues remain open. First, we have focused on the "2-scheme" regime (where 3g ) in order to track Adam's 's preconditioner over a long timescale; extending our analysis to the intermediate 1.5-scheme or other scalings of 1 - 33 is left for future work. Second, our derivations assume that the iterates remain close to a smooth minimizer manifold; understanding Adam's implicit bias once the trajectory ventures beyond this local neighborhood may require restarting the analysis from the SGD dynamics. Finally, our approach cannot cover weight-decay or decoupled decay terms such as the W-term in Adam W; characterizing how weight decay alters the effective sharpness regularizer is an important direction for follow-on studies.
6 Acknowledgement
We would like thank Xinran Gu, Kaiyue Wen, Yushun Zhang, Jiaye Teng, Yuhang Cai, and the anonymous NeurlPS reviewers for their insightful comments and feedback.
10

---

References
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprini urXiv:2002.09011 2020.
Vaswani Ashish. Attention all you need. Advances in neural information processing systems, 30:1, 2017.
David GT Barrett and Benoit Dherin. Implicit gradient regularization. arXiv preprint arXiv:2009. //660 2020.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an crnttin-uulenbeer like process. In Conference learning theory, pages 483-513, PMLR. 2020.
David Brandfonbrener and Joan Bruna. Geometric insights into the convergence of nonlinear td learning. arXiv preprint arXiv:1905.1211: 12185, 2019.
Matias D. Cattaneo and Boris Shigida. How memory in optimization algorithms implicitly modifies the loss, abs//baaa////// 2502. 02132
2025. Utt.htttee//arxxv..oo wrtpa///ww....r.aa//0...22.
Matias D. Cattaneo, Jason M. Klusowski, and Boris Shigida. On the implicit bias adam, 2024. URL https://arxiv. http:////aaxiv.oooaaaab///////// 0009
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Alex Damian, Tengyu Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449-27461 2021.
Alexandre Défossez, Léon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. arXiv preprint arXiv:2003.00 222220003.29,,,.... 2020.
Steffen Dereich and Arnulf Jentzen. Convergence rates for the adam optimizer. arXiv preprint 2rXiv:2407.21007, 2024.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 4171-4186, Min- neapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 11.18663/1//11914400 URL https://aclaattoolud wtg.
Alexey Dasovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbom, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.11929, 2020.
Aijun Du and Jinqiao Duan. Invariant manifold reduction for stochastic dynamical systems. arXiv preprint math/0607366, 2006.
Johannes Jisse Duistermaat and Johan AC Kolk. Lie groups. Springer Science & Business Media, 2012.
KJ Falconer. Differentiation of the limit mapping in dynamical system. Journal of the London Mathematical Society, 2(2):366-372, 1983.
Benjamin Fehrman, Benjamin Gess, and Amulf Jentzen. Convergence rates for the stochastic gradient descent method for non-convex objective functions. Journal of Machine Learning Research, 21(136):1-48, 2020.
Damir Filipović. Invariant manifolds for weak solutions to stochastic equations. Probability theory and related fields, 118(3):323-33) 341, 2000.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021. URL htttp://openzerierr wttp://oroommmdddddmm.mmmmoo
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry Vetrov, and Andrew Gordon Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns, 2018. URL httpsttss:/// / //arxvv.///ii//aaaa org/ abs/182...1 10026
Khashayar Gatmiry, Zhiyuan Li, Tengyu Ma, Sashank Reddi, Stefanie Jegelka, and Ching- Yao Chuang. What is the inductive bias of flatness regularization? study of deep matrix factorization models. Advances in Neural Information Processing Systems, 22288000220555 2023.

---

Jonas Geiping. Micah Goldblum. Phillip E. Pope. Michael Moeller, and Tom Goldstein. Stochastic training is not necessary for generalization, 2022. URL httpttt///a/iii..oo wwwpiiivorg//aba haa///aa///// 2209 14119
Xinran Gu, Kaifeng Lyu, Longbo Huang, and Sanjeev Arora. Why (and when) does local sgd generalize better than sgd? arXiv preprint arXiiv2303.0121.. 2023a.
Xinran Gu, Kaifeng Lyu, Longbo Huang. and Sanjeev Arora Why (and when) does local sgd generalize better than sgd?, 2023b. URL httpp://arxiv org/ abs 200.0012153100
Xinran Gu, Kaifeng Lyu. Sanjeev Arora, Jingzhao Zhang, and Longbo Huang. A quadratic synchronization rule for distributed deep learning. In The Twelfih International Conference Learning Representations, 2024.
URL U http://e//ppeeevieeeeeeeeeeeeeee
Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianbao Yang. Unified convergence analysis for adaptive optimization with moving average estimator, 2025. URL
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842-1850. PMLR, 2018.
Jeff Z. HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Conference on Learning Theory, pages 2315-2357. PMLR, 2021.
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Cited on, 14(8):2, 2012.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1::1-42, 01 1997. ISSN 0899-7667, doi: 10.1 10.1162/neco.1999 1..11. /ttp::/ts://do. wtt.:/.o//01..../eeeeeeee.e http://oeco.999.....0 .9.9799...1
Yusu Hong and Junhong I. Lin. High probability convergence of adam under unbounded gradients and affine variance noise, 2023. URI httaappe/aarra...oo/aaa//211111.
Hideaki liduka. Theoretical analysis of adam using hyperparameters close to one without lipschitz smoothness, 2022. httpss/:///axx org/ absp///.....22..///// 2206. 13290
Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexci A. Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE Conference Computer Vision and Puttern Recognition (CVPR), pages 1125-1134. IEEE, 2017.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemnl:: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima sgd. arXiv preprint arXiv: 17110063, 2017.
Ziwei Ji and Matus Telgarsky. Directional convergence and alignment deep learning. In H. Larochelle, M. Ranzato, R. Hadsell. M.F. Balcan, and Lin, editors, Advances in Neural Information Processing Systems. volume 33, pages 17176-17186. Curran Associates, Inc., 2020. URL URt httpps:/pproeeddii neurips cc/paper_ files/ httpa////ppeee//0///////////6//a
19w5c0d11122e0eee55 Pwper. .ddf
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishman, and Samy Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2020. URL.
https://openreevev net. wwtp/eeoweiddds........
Ryo Karakida, Tomoumi Takase, Tomohiro Hayase, and Kazuki Osawa. Understanding gradient regularization in deep learning: Efficient finite-difference computation and implicit bias, 2023. URL htttp:///aax:: org/ abb//21b..2.....///
2210 .0720
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017. URL J |tttps://openreviess net/ forum?id-H1oyRYy
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 142.6:1412..9980. 2014.
Rohith Kuditipudi, Xiang Wang. Holden Lee, Yi Zhang, Zhiyuan Li, Wei Hu, Rong Ge, and Sanjeev Arora. Explaining landscape connectivity of low-cost solutions for multilayer nets. In Wallach, H. Larochelle, A. Beygelzimer. F. d' Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing wea1i99.RRLteee UtL h/ttttpppp//prroeeiing.... nourips. cct.ccccppaaa ccppaper Paper sopprr.pfff
Systems, volume 32, Curran Associates, Inc 2019,
12

---

J. Lee. Introduction to Smooth Manifolds. Graduate Texts in Mathematics. Springe New York, 2012. ISBN 9781441999825. URL httptttt///::bkss wwwpllkss.oooo..o. www.g/o.....././aaeo ptt.g//......ooooaaaaaaoggc./..a
Qianxiao Li, Cheng Tai, et al. Dynamics of stochastic gradient algorithms. CoRR, 2015.
Qianxiao Li, Cheng Tai, et al. Stochastic modified equations and adaptive stochastic gradient algorithms. In
International Conference on Machine Learning, pages 2101-2110. PMLR, 2017.
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and dynamics of stochastic gradient
algorithms i: Mathematical foundations, 2018. URL https: http:///axxvvvvv. www.iv..org/aaeo hbsgg//b///// abs/ //////811110001 01558
Qianxiao Li, Cheng Tai, et al. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research, 20(40): 147, 2019.
Zhiyuan Li. Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochastic differential equations (sdes). In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 12712 12725. Curran Associates, Inc., 2021a. URL https::/ pdf
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss?-a mathematical framework. arXiv preprint arXiv:2 10.02/10000694,, 2021b.
Tao Lin, Schastian U Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini-batches, use local sgd. arXiv preprint / 000.0882777777 2018.
Wu Lin, Scott Lowe, Folix Dangel, Runa Eschenhagen, Zikun Xu, and Roger B Grosse. Understanding and improving the shampoo optimizer via kullback-leible minimization. arXiv preprint arXiv:2509.0333 2025.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Leurning. volume 202 of Proceedings of Machine Learning Research, pages 22188-22214. PMLR, 23-29 Jul 2023. httLhttp://// / http:////oceeedeeneeeeee mlr. press/v202/ http://222/11uu2aa 1iu23ao...... html
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, arXiv preprint vrXiv:1711.550000 2017.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference Learning Representations, 2020. JRLLAatps///opoeeiiiee.eeee////e gd=SJeLI
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. Advances in Neural Information Processing Systems, 35:34689-34708, 2022.
Kaifeng Lyu, Jikai Jin, Zhiyuan Li, Simon S Du, Jason D Lee, and Wei Hu. Dichotomy of early and late phase implicit biases can provahly induce grokking. In 12th International Conference Learning Representations, ICLR 2024, 2024.
Chao Ma. Daniel Kunin, and Lexing Ying. A quasistatic derivation of optimization algorithmss' exploration on minima manifolds, 2023. URL/tttps: httttpps:///pnnreviiee net. tete:/foorumidddDbbbbt -3
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for
adaptive gradient algorithms.. Advances Neural Information Processing Systems, 35:7697-7711, 2022.
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms, 2024.
10287
Volodymyr Mnih Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533 2015.
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on shampoo's preconditioner. arXiv preprint 2xXi:2206.11774. 2024.
Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry. Lexicographic and depth-sensitive margins in homogeneous and non non-homogeneous deep models. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 4683 4692. PMLR, 09 15 Jun 2019.
13

---

Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.6614, 2014.
Zachary Novack, Simran Kaur, Tanya Marwah, Saurabh Garg, and Zachary C. Lipton. Disentangling the mechanisms behind implicit regularization in sgd, 2022. URI.http:: ptttppa://arzii http//aawvvv//aaaaaaeaeea aaaaa2aa11221 2211. 15853
Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
Qian Qian and Xiaoyuan Qian. The implicit bias adagrad on separable data, Advances Neural Information Processing Systems, 32, 2019.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. Unpublished manuscript, 2018.
Sashank Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference on Learning Representatons, 2018.
Olaf Ronneberger. Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention (MICCAI), volume 9351 of Lecture Notes Computer Science, pages 234 241. Springer, 2015.
Nikunj Saunshi, Stefani Karp, Shankar Krishnan, Sobhan Miryoosefi, Sashank J. Reddi, and Sanjiv Kumar. On the inductive bias of stacking towards improving reasoning. In Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang, editors, Advances Neural Information Processing Systems, volume 37, pages 71437-71464, Curran Associates, Inc., 2024. URL https: httpttt:///occeedingggeee////// neur wwtpeiissccppppppeee cc/paper_ files/ 2llle///pperr2224//eea http://///.///0ccaaaa22222222223
Paper Conference pdf
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization
algorithms. arXiv preprint arXiv: /00::/777.077777
2017. 2017.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604, PMLR, 2018
Hao-Jun Michael Shi, Tsung Tsung-Hsien Lee, Shintaro Iwasaki, Jose Callego-Posadee Zhijing Li, Kaushik Rangadu-
rai, Dheevatsa Mudigere, and Michael Rahbat. distributed data-parallel pytorch implementation of the distributed shampoo optimizer for training neural networks at-scale. arXiv preprint orxi:22309.06/3 2023.
Naichen Shi and Dawei Li. Rmsprop converges with proper hyperparameter. International conference on learning representation, 2021.
Samuel Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 9058-9067. PMLR, 13-18 Jul 2020. httpttt::///oceeeeeeeeeee
mlr. wttps//vee//e/ppppp..p. wwwp//aa.tttaa.mmmmmm html
Samuel L. Smith, Benoit Dherin, David G. T. Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent, 2021. URL https://aaax org http::////azv...rr/aaa////..111
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasckar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1-57, 2018
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Bohan Wang, Qi Meng, Wei Chen, and Tic- Yan Liu. The implicit bias for adaptive optimization algorithms on homogeneous neural networks. In International Conference on Machine Learning, pages 10849-10858. PMLR, 2021.
Bohan Wang. Huishuai Zhang, Qi Meng. Ruoyu Sun. Zhi-Ming Ma, and Wei Chen. On the convergence of adam under non-uniform smoothness: Separability from sgdm and beyond, 2024a. URL httpp://aaxxv org/ ab////aa/40 15146
Bohan Wang. Yushun Zhang, Huishuai Zhang, Meng, Ruoyu Sun, Zhi-Ming Ma. Tie- Liu, Zhi-Quan Luo, and Wei Chen. Provable adaptivity adam under non-uniform smoothness. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD '24, page 2960 2969, New York. NY, USA. 2024b. Association for Computing Machinery. ISBN 9798400704901. doi: 10.1145/36352283377177
URL https: / /doi. org/10. 1145/3637522 3671718
14

---

Jiaxuan Wang and Jenna Wiens. arXiv:206.1654), 2020.
Adasgd: Bridging the gap between sgd and adam. arXiv preprint
Runzhe Wang. Sadhika Malladi, Tianhao Wang, Kaifeng Lyu, and Zhiyuan Li. The marginal value of momentum for small learning rate sgd. arXiv preprint arXiv:2307155190 2023.
Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, and Tengyu Ma. Understanding warmup- stable-decay learning rates: A river valley loss landscape perspective, 2024. htt.attta://arriv.co org/
abss/2110.05
Blake Woodworth. Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes overparametrized models. In Conference on Learning Theory, pages 3635-3673. PMLR, 2020.
Shuo Xie and Zhiyuan Li. Implicit bias AdamW: fox -nom constrained optimization. Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference Machine Learning. volume 235 of Proceedings of Machine Learning Research, pages 14488-54510. PMLR, 21-27 Jul 2024.
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa Bianchi, and Garnett, editors, Advances Neural Information Processing Systems, volume 31. Curran sociates, Inc., 2018. URI. httpps:/pproeeedingg www.gg......r........ wwtprppp..cppppprrfiiiiii.e http:/ppppe..leeeepp/ee///////// 0036555ccc7477333004ceeeaa222-p2 pdf
Chenyang Zhang, Difan Zou, and Yuan Cao. The implicit bias of adam on separable data, 2024a. URL
httpa://arxiy: http:///xxvv.v.. www.aa..gwwb.aaaa..a.orrraaa..aa www.//gg///bbeeeeeeeo htrg//bb/24006 1065/
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning
requires rethinking generalization. In International Conference on Learning Representations, 2017. URL
htttps:/ooperreviet www. weee//oorumidd-yygeeee
Qi Zhang, Yi Zhou, and Shaofeng Zou. Convergence guarantees for rmsprop and adam in generalized-smooth non-convex optimization with affine noise variance. arXiv preprint arXiv:2404.0145: 2024b.
Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. Adam can converge without any modification update rules. Advonces neural information processing systems, 35:28386-2888 2022.
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu L., Ruoyu Sun, and Zhi-Quan Luo. Why transformers need adam: A hessian perspective, 2024c. (ttp::ttts://///// //wp:i/.axxivvv//a httppa////xxvvoor//aaaaaa www///////////////////.//....a/a wttaaaa/..22222 2202..16888.888
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P. Kingma, Yinyu Ye, Zhi-
Quan Luo. and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more, 2025. URL https // /aaxii..org wwrriv.org/aba... 2406 16793
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade. Deconstructing what makes good optimizer for language models, 2025. URL httttss:/aaai..org///////
.07972
Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization in deep learning, 2022. URL /tttt:t//axxiii/ wttp:v/..rrgabaa/////// 03599
Fangyu Zou, Li Shen, Zequn Jic, Weizhong Zhang, and Wei Liu. A sufficient condition for convergences of adam and rmsprop. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pages 11127-11135, 2019.
15

---

1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer:
[Yes]
Justification: Our contributions are accurately summarized in abstract and introduction, with no overclaims.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contri- butions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
See the discussion section.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-tex, -o-tttxttttxttt system might not be used reliably provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems privacy and faireess.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer:
[Yes]
Justification:
We provide the proof for each theoretical result in either main part the Appendix.
16

---

Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can cither appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: we provide details and concrete settings of our experiments the experiment sections.
Guidelines:
The answer NA means that the paper does not include experiments. the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, may be necessary to either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility how to replicate the results, access to hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, other means that appropriate the research performed.
also be provided via detailed instructions for
While NeurIPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reprodueibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material? Answer: [Yes]
17

---

Justification: We provide details for the dataset construction and instructions for reproducing our experiments the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (attps://nips.cccppbbbl)
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/
Seaiiddes/CodeSSubiissiiiii
for more details.
The authors should provide instructions on data access and preparation, including how to access the data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
experimental results for the new
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparame- ters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide the optimizers and main settings the experiment section, and other details in the supplemental material.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification:
We report error bars across multiple random seeds for all main experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main
claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar the standard deviation or the standard error the mean.
18

---

It is OK to report 1- 1-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have a 96% CI, if the hypothesis Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are of range (e.g. negative error rates). error bars are reported tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Answer: [Yes]
Justification:
We detail the total compute resources used in Appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Question:
Answer:
[Yes]
Justification:
We have read and complied with the NeurlPS ethics review guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there is special consideration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), faireess considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
19

---

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mecha- nisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer:
[Yes]
Justification:
We cite the creators of all external datasets and assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi thcode. com/ hommdatasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification:
[NA]
Guidelines:
The answer NA means that the paper does not release assets.
20

---

Researchers should communicate the details of the dataset/code/model part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
21

---

Contents
Introduction
2
Related Work
Preliminaries
4
Theoretical Analysis of Adam
4.1
Slow SDE Analysis for AGMs
Interpretation of The Slow SDEs for AGMs
Technical Difficulties and Proof Insights
9
4.3.1
Convergence Guarantee of AGMs
4.3.2
Key Insights in the Derivation of Slow SDHs for AGMs
10
5
Conclusions
10
Acknowledgement
10
Illustration of the Difference between Conventional SDE and Slow SDE
24
Adam's Provable Generalization Benefit with Label Noise
24
B.1
Implicit Regularizer of Adam under Label Noise
24
B.2
Example: Sparse Linear Regression with Diagonal Net
26
B.2.1
Result: Adam's Implicit Regularizer Facilitates Sparse Ground-trunt Recovery
26
Matrix Factorization: Adam Implicitly Regularizes Sharpness Differently
27
Problem setup
27
Results
28
D Restatements of the Main Results
29
D.1 Slow SDE for AGMs
29
D.2 Convergence Guarantee AGMS
30
Constructing the Working Zones
31
Proof of the Convergence of AGMs
33
F.1
Proof of Saaverrgnce-Relltte. Conclusions in Appendix D
40
Proof of the SDE Approximation AGMs
43
Lemmas for Adaptive Manifold Projection
43
G.2
Iteration Stays Near Manifold
45
Moment Calculation of AGMs Near Manifold
47
G.3.1
Moment Calculation Within a Giant Step
50
G.4
Weak Approximation
55
22

---

G.4.1
Preliminary and Additional Notations
56
G.4.2
Proof of the Approximation for Slow SDE of AGMS
............
57
G.5
Proof of Theorem G.2
59
H
Proof of T'eorems in Appendix B
59
H.1 Proof of Adam and AdamE's Implicit Biases with Label Noise
................................
59
Proof of Lemma B.3
62
Regularizers under label noise for AGMs in Table 1
63
Shampoo Optimizer as an AGM
64
U.1
A brief introduction to Shampoo
64
J.2
Shampoo under the AGM Framework
65
U.3
Vectorized AGM form
66
J.4
Discussion on Shampoo's ' Implicit Bias under Label Noise
66
23

---

A Illustration of the Difference between Conventional SDE and Slow SDE
In this section, we illustrate the difference between conventional SDE and slow SDE. In Figure 2 let I' denotes a ID manifold, then the discrete iteration of the optimization process can be seen as successive steps (orange, Figure 2ah that starts from A, first converge to some point B in and then move along I to C.
A
B lteration
Conventional SDE
H
B
Slow SDE
r
(a)
(b)
(c)
Figure 2: Comparison of conventional SDE and slow SDE.
The main intuition behind slow SDE is that the whole process A + B decomposed into two motions: convergence motion A - H (dashed, Figure 2c) and an implicit regularization motion H + B + C. The convergence motion is fast and dominates the dymanics convergence phase ends; meanwhile the
can actually be
during the convergence phase, but it fades out slow, implicit regularization motion starts to dominate.
The conventional SDE approximates the convergence phase only, whose unit time corresponds to O1) steps Figure 2b In contrast, slow SDE manages to separate the slow implicit regularization motion from the fast convergence, and approximate the implicit regularization near manifold only Figuuee 2c).
B - C) varies in the analysis of different optimizers. Intuitively, the projection should reflect the converging direction driven by clean (without noise) and continuous version of the optimizer. In SGD the projection is gradient flow; but in Adam we need to consider the preconditioning effect caused by 1/ (Vu + e), so add an SDE to track the preconditioner, and define preconditioned gradient flow for projection.
Remark A.1. The projection method (which projects A + B Co
B Adam's Provable Generalization Benefit with Label Noise
In this section, prove that with the label noise condition, the implicit regularizer of Adam reduces to simpler form that aligns better with sparsity regularizations, and then verify experimentally.
B.1 Implicit Regularizer of Adam under Label Noise
Label Noise. By label noise we refer to the condition that for all E: r, the covariance matrix Σ is constant multiple of the Hessiam: E())
((C(0) for some constant a [Blane et al. 22020). This condition typically arises when training a model that is overparameterized enough to the training data perfectly, while fresh noise is added to the target label at each step of training.
To see how this label noise condition holds, imagine a simple regression problem with training data {(x(f) ( } and model h(0: æ). The loss function for a single data sample is defined as ((0: mr, v) (h((0; y) With label noise, at each step of training, data sample (x,, 3t) as well as an independent noise perturbation & from distribution with zero mean and constant variance 8 0. Then we take one gradient step on ((0; ; æt, Yt (e) with the noisy label yt + L..
randomly take random
The expected training loss becomes (()) Ej æt, Ur 4.).. When the model is sufficiently overparameterized so that h(0 x) can be simultaneously attained for all training data points in its parameter space, the minimizer manifold just r {0 L(0) = 1 _ where _ entirely due to the presence of label noise. On this manifold, we can how the label noise condition
24

---

00 a (0) holds for a = =
0 : E Ve(0; æz, yr + ( Mr +G
E = ~V
Beyond this simple setting, similar analysis can be applied to establish the label noise condition (0) 2² for mini-batch training and other common losses such as cross-entropy loss. This proportional relationship between Hessian and noise covariance greatly simplifies the analysis of training dynamics and has been widely used to study the implicit bias of SGD and related optimizers
Blanc et al. 2020 Damian et al. 2021 Li al. 2021b Gu al. 2023a
Under the label noise condition, Li ct al. 20216 proved that slow SDE for SGD reduces to an ODE. In the following theorem, we show that the slow SDE for AGM also reduces ODE, but the resulting ODE takes a different form:
Theorem B.1 (Slow ODE for AGMs with Label Noise). Under the label noise condition, the Slow SDE AGMs in Equation (2) becomes the following ODE:
{
dc(l) - - dv(t) c (VE(()) v) dt,
-,S.0 ( C))
where S, : S(v()))
See Appendix for the proof. We then derive the implicit bias of Adam with label noise.
Lemma B.1 (Adam's Implicit Bias under Label Noise). Under the label noise condition, every fixed point of Equation (3) for Adam satisfies Vrtr (Diag(H) (2)) O.
Proof Sketch. At the fixed point of Equation (3) we have
== V(E(C)) S(v)00 (C0p) = =.
With label noise, H def V²C(() = E(C)/a. In the Adam case, = diag Σ) = odiag( (H), and
S(v) Diag(1/ /v) with e 0. Then S(v) Diag( (adiag H)) ) and we can simplify the regularizer term into
3 [(C)(S(())]
= Σ j=1 (aH,,) V(H,))
2 Vir (Ding(H)) ((HH)))) = Va
Va
j=1
which implies that the stationary point of Adam satisfies
S(v) s())) Vtr (Ha H
0.
The preceding matrices have some clean properties. From Lemma G.2 0 (C) acts as an invertible linear map on the tangent space of the manifold I, and vanishes on its normal space. Together with the invertibility of S(v), this gives
Vrtr (iag(H)) 1/2)
for any fixed point of Adam's 's slow ODE. See Appendix H Horr more details.
The above proof can be generalized the case of 0. When e >> the implicit bias of Adam changes to
Diag( 1/2 c (
Vrtr
= 0.
Va
See Lemma H.2 for more details. Note that this trace term is always non-negative, since x ya for any 2 0.
r + 2 Vo c
va
Simple Way Tune Adam's Implicit Bias: AdamE. The proof of Lemma B.I inspires the following simple variant of Adam: We define AdamE as an optimizer class that, is identical to Adam except that S(v) DDaag11((22)) + e)) for tunable parameter E [o, 1). For any X E [0, 1) we also use the term AdamE with À = λg. simply AdamE-Ao. Note that AdamE with À = coincides with Adam, and that AdamE optimizers lie within the AGM framework. Applying the same method as in Lemma B.I yields the implicit bias of AdamE under label noise.
25

---

Lemma B.2 (AdamE's Implicit Bias with Label Noise). Under the label noise condition, the fixed point of Equation (3) for AdamE-A with [0, 1) and = 0 satisfies Vrtr (Diag(H)'>)
Lemma B.2 indicates that tuning the exponent of the second-order moment in Adam exactly results in tuning the exponent of diag( V [()) in the implicit bias. When À = 0. the implicit bias reduces to that of SGD, and AdamE also gets rid the effect of second-order moments and reduces to SGD with momentum, which coincides perfectly. Next, we relate the implicit bias to sparsity and compare the performance of Adam, AdamE, and SGD a simple experimental setup.
B.2 Example: Sparse Linear Regression with Diagonal Net
In this section, we adopt the diagonal linear network (diagonal net) setting proposed by Woodworth et al. [2020] an experimental setting, which is also used by Li et al. [202bb] to study the implicit bias of SGD.
Setting (Diagonal Net with Label Noise): Let w" E R be an unknown s-sparse ground truth vector.
Let {(z;, y:) se)] be the training dataset where each 2i ~
[ [=1]". and each yi is generated by (zt, w*) Our parameter is defined O (#) TR 2d For any function g defined on R2d, we write
Lid. Unif
g(0) and g(u, v) interchangeably. The loss function is defined as:
(()) - (0), where (, 2 ((2,, u2 0 vi)² n
1
where a label noise is added to the true label y during training. This setting can be viewed as using estimation w u 02 to approximate the ground truth vector w² of linear regression task. Note that d > n here so the model is highly overparameterized: Theoretically, Li et al. [0021b] proved that experimentally that less than 1000 training pairs is required for both Adam and SGD to achieve a low test loss when d = 10000. The manifold is defined wherever zero train loss is achieved, i.e. r= {0] (z, u 0 ys. Vi E [rl]].
O(x In d) is enough for SGD to recover ground truth, and will later show
This setting allows to relate the implicit bias directly to the sparsity of the estimated ground truth. Lemma B.3. Let 0 be an optimal parameter minimizing the loss function C. i.e. 0" For each
0 = (*) E r, denote w == 0 a 02 and H := V V'C( We have the following:
ff E arg tr(iiag (H)..)) then we also have 0" E aar arrmmnnger |/0l...5
Furthermore, for any ep E (0, if 0" E arg minger ([Diigg()))) then we also have 0' E arg migger |
The main idea of the proof is that the training loss depends only the combined quantity w 0 o2 Hence, if for some index both u; and iÚ are nonzero, we can reduce the magnitudes of uq and U, while keeping u? of fixed, obtaining another minimizer with strictly smaller tr(Diag H ).)) Therefore, at any optimum we must have u, = 0 or v, = O for every i. Under this condition, ((Diag(H)*) can be identified with || We provide the detailed derivation in Appendix H
Lemma B.3 gives the following insights: Implicitly regularizing tr (Digg H)) is equivalent to regularizing the -nom of the estimated ground truth wü = 0 v 02 Adam corresponds to fo.5. SGD to L, and AdamE-A to 41-x. Just as lasso (l1) is preferable to ridge (62) for sparse ground- ground-truth recovery, we therefore expect Adam and AdamE (with À > 0) to recover sparse ground truth more efficiently than SGD. We verify this prediction below.
B.2.1 Result: Adam's Implicit Regularizer Facilitates Sparse Ground-truth Recovery
Figure 3 shows the results of the experiment. We gradually increase the number of training points and train Adam, SGD, and AdamE under several configurations until convergence. We consider a configuration to have recovered the ground truth if the test loss falls below As illustrated in Figure 3a Adam's 's test loss plunges towards zero at approximately Purain = 420, whereas SGD's test loss decreases more gradually as the training set grows. To interpolate between different implicit biases, we evaluate AdamE for several values of X. Figure 3b indicates that AdamE, even with small positive value of X, exhibits the same sudden recovery behavior as Adam. This suggests that Adam's 's influence on the implicit bias upon SGD arises from its preconditioning mechanism.
26

---

Adam vs. SGD
AdamE vs. Adam and SGD
Rinal tast loss taymigg saae,, throshod-11
10
100
Adme 001
Adam 01
MamE 0
Adanã 0.75
AdanE 09
Adam
Adam
SGD
300
400
500
600
700
800
300
400
500
600
700
800
(b)
Figure 3: Final test loss function of the training data size with d = 10000, 50. Each plotted point is the final test loss after both the training and test losses have converged; its x-coordinate is the training data size and the curve denotes the optimizer and configuration. (a) Loss comparison between SGD with different learning rates, and Adam with different learning rates and 82 values. (b) Loss comparison between AdamE with À 0.01, 0.1, 0.25, 0.75, 0.9, Adam, and SGD.
Takeaway. Adam's unique implicit bias may help to recover sparse ground truth.
However, we should also keep in mind that a clear interpretation of Adam's unique implicit bias, tr(Diag 1/2) relies heavily on the condition that H is diagonal. Only with this condition can we claim Adam as minimizing [Auloo.s instead of SGD's |//l where h is the vector consisting of all cigenvalues of H. In other words, Adam's optimization on the implicit bias upon SGD only makes sense when H is diagonal. In the diagonal net setting this is indeed the case in expectation, but will see the next chapter that Adam's 's unique implicit bias may even lead worse generalization when H is no longer diagonal.
C Matrix Factorization: Adam Implicitly Regularizes Sharpness Differently
The diagonal experiments in Section showed that Adam's implicit bias towards sparsity improves generalization relative to SGD. We now turn to supply the potentially negative impact of Adam's implicit bias in another controlled setting: deep matrix factorization with label noise, where the relevant implicit regularizers are analytically tractable.
In this task, Adam is expected to minimize tr(Diag HH2 rather than (rHH) Leveraging existing theory, we therefore predict that (i) Adam will converge to solution with larger tr (r(H), whose tr(Diag( H)//2) however, smaller than SGD's solution, and (ii) once Adam reaches a solution with larger tr H), supported by Gatmiry et al. [003)), will generalize worse than vanilla SGD in the presence of label noise. Our experiments confirm both predictions (Figure 4add Figure 5).
C.1 Problem setup
We consider an L-layer linear network with parameters W = (W1.... WL), where W, Eddxd. and d; (do, dr.) for all We M dR d be a rank-r ground truth matrix and observe the n i.i.d. linear measurements ((A;, ) generated by b, (A, M With label noise and mini-batch size B, the empirical loss at step is:
((W
Σ ((A,, W. W1) b, + Ee,))
B 1Eb,
where B, stands for a fresh batch of size B, and the label noise Gr,~ o2) are independent to each other for both different t and i.
27

---

Depth Hessian and Loss Metrics
Trace of Hessian
Trace of sqrt Hessian
SGD Adam
SGD
Trace (log]
Trace (log)
Adam
10a
10a
10a
10
10
10
10a
10
10
10y
10
10
10
10a
10
10
10'
(log)
Iteration (log)
Training Loss
Test Loss
10a
MSE Loss (log)
SGD
MSE Loss (log)
SGD
10F
104
10-3
10-2
10a
10
102
10
10
10
10
10'
10"
10
102
10
10 (log)
10
10"
10
Iteration (log)
Iteration
Figure 4: Deep matrix factorization with label noise with deepth L = 2. Adam and SGD are trained on identical data and noise realizations. Top: evolution of tr( H) and tr( H)1/2). Bottom: training and test MSE. Adam converges to a point with larger tr( H) but smaller tr(Diag( HH) and exhibits higher test error.
Depth 5: Hessian and Loss Metrics
Trace of Hessian
Trace sqrt Hessian
Trace (log)
20a
Trace (log)
10a
10a
SGD
Adam
10
10a
10
10a
10a
10
10
10c
107
100
10a
10a
10a
10
10f
107
Iteration (log)
Iteration (log)
Training
4 10"
MSE Loss (log)
10F
SGD
MSE Loss (log)
SGD
Adam
3x 100
Adam
10-3
2x
10-1
10a
10'
10
10a
10
102
10c
10y
10a
10a
102
10a
10f
10a
10f
107
Iteration (log)
Iteration (log)
Figure Deep matrix l'actorization with label noise with deepth L 5.
C.2 Results
Our setups for SGD optimizer follows Section of Gatmiry et al. [2023 For Adam, we use the standard hyperparameters 81 0.9, $2 = 0.999, and learning rate 10 3. all other settings are identical SGD.
Figure 4 and Figure 5 (top rows in the figures) show the evolution of sharpness metrics and the training/test losses for layer depth L = 2, Adam drives tr(Diag 1/2) ()) sharply downward while tr(H) remains high and even non-monotone, confirming that Adam does not target Hessian trace. Correspondingly, the bottom rows in these two figures show that Adam attains a higher test loss despite similar training error, which is evidence that Adam's 's implicit bias is detrimental in this setting.
Recall that SGD with label noise implicitly regularizes ir( H) on the minimizer manifold. On this matrix factorization task, Gatmiry et al. [2023) proved that minimizing ir(H) roughly equivalent to minimizing the nuclear norm of W which is connected to generalization improvement when M*
28

---

possesses the low-rank property. In contrast, our theoretical analysis suggests that Adam implicitly regularizes tr (r(diag( ( rather than tr( H)), inducing a different bias that does not necessarily favor low-rank solutions. Consequently, Adam is expected converge to solutions with smaller tr(diag( ) larger (r(H) and worse generalization compared to SGD.
Takeaway. deep matrix factorization with label noise, Adam's tendency to minimize (H)iaa(())/)) drives it towards solutions different from those found by SGD and exhibits worse recovery of the low-rank ground truth. This highlights that Adam's 's implicit regularization can be detrimental for certain tasks.
D Restatements of the Main Results
In this section, to make results more organized and easy to understand, give the restatements of the main results in Section 4] where we present the two main theorems.
1. The AGM iterates converge neighborhood of the manifold (Theorem 4.2
2. Morcover, once the iterates enter this neighborhood, their dynamics over On1 ?) discrete steps can be accurately tracked by a slow SDE Theorem 4.1).
Recall that in the AGM framework, the transition from 0₂ to 0x+1 is defined as:
mg+1 == 8 m. + (1 (
0k+1 82 (1 (V)
Ok+1 O2 118( 115(0)+1)
under the following conditions:
1. S TR - $d is p.-smooth, where TR denotes the subset of TR. D that has non-negative entries, and - denotes the space of dx positive definite matrices.
2. ((v) Ho I for some Rg and any e E R.D 200
3. V _dxd
RDD is linear, and (ggg T) E R for all g E Rd.
D.1 Slow SDE for AGMs
Theorem D.1. Let Assumptions 3.2 3.3 and 4.2 be satisfied. Let denote a local minimizer manifold,
n be sufficiently small learning rate of an AGM, and B2 (g7)). Then we have the following conclusions:
1. (Convergence to a near-manifold neighborhood) There exists a constant co > such that for any initial point O, whose &z distance from I does not exceed cu, and any 8 E (17200 1), with probability at least 1 the following holds for some Ko ( log 1):
L(OKe) -L- -L'=0 -===0(100s
1
n6
0K - $sK, (0xo)ll (0 (v(nnoo n5
2. (Slow SDE tracks AGM's 's trajectory in weak approximation sense) Moreover, let T K == [Tn 2), and suppose Assumptions 3.4 3.1 and h.ld.... Continue rumning an AGM for K iterations after reaching the final state (0K.. UKa conclusion I and totally obtain ((0x, Ko+K Let x(t) = (C(t), v(t)) be the solution 2 with initial condition
x(0) (x)=Xxe0 E IR
X ( (0ko). vx.)
'The exponent here, along with the exponents related to the 5-goodness in Definition G.1 can be arbitrary large constant, which does not affect the order of following derivations.
29

---

For any k E (0,1..... J.. let Xx = ( ((s (KKo+))) (Ko+k) be the AGM state at step Ko + k projected onto r. Then for any C function g
max
[ xx)
X
[g (oy ( = Xo] =
OSKSK
where Ö()) hides factit and constants independent of (but possibly depending g).
Notice that in Theorem D. we give more explicit statement for our SDE approximation results in Theorem 4.1 One can see that Theorem i. direct corollary of Theorem D.I
Proof of Theorem D.7 For the convergence part of Theorem D.I we first construct the working zones near the manifold in Appendix to ensure some properties that are crucial to our analysis such as the j- condition. After that, we give the proof of convergence in Appenddix F. We finish our proof Appendix by combining the previous results in working zone and convergence analysis.
And the proof of the SDE approximation part is finished in Appendix G.4 through the moment calculation and the subsequent weak approximation analysis in Appendix G
D.2 Convergence Guarantee of AGMs
In the proof, the first part of Theorem D.I 1.1pp is done by first proving a convergence result with global μ-PL condition, and then arguing that AGM starting near enough to the manifold will stick to the manifold with high probability. As mentioned in [Sectiinn 4.3.I] the convergence under μ-PL condition be
separate technical contribution of our paper, which is stated below.
Definition D.1 (Polyal--.ojasiewwicz Condition). For some μ., L 0, we say some function L Rd - R is (M, Poooyak-Lojasievii (abbreviated as (n, L)-PL) if and only if for all O E IRd such that ((0) S L:
L*) 2p(((()) - < |
When +00, we call this condition the j--Pllyk- (μ-PL) condition.
Theorem D.2 (Formal restatement of Theorem 4.2). Let Assumptions 3.2 3.3 and 4.2 be satisfied, C
be function sattisfyngg the μ-PL condition, and 1) be sufficiently small learning rate of an AGM. Let 1- 32 = O(7²) For any e (0, 1), with probability at least 6, the following holds for some K 0 log
[(Ok) - LL' -L'=0 1100 //
-L'=0(nnop
1
176
10x Φ5 (0x)l12 (V no
Note that Theorem is different from Part 1 Theorem D.I in the sense that Theo- rem requiies the p-PL condition to hold globally, while Part heorem D. not, Actually, will find from Lemma E.3 that p-PL provably exists in a neighborhood of r, and we prove that an iteration of AGMs starting within that neighborhood stays within that neighborhood with high probability.
Remark.
the latter requires the iteration start from some neighborhood T. Later
There have been many previous works discussing the convergence bound of Adam. For example, Reddi et al. (2018) and Dereich and Jentzen (2024) give convergence bounds under the convexity
condition, Zou et al. [2019]. Shi and Li 202 and Zhang et al. [2022 focus on the cases where
learning rates follow a 1/vl f decay, and the bounds given by Zaheer al. 20188) Zhang et 2022
and Wang et al. 2024b) do not decrease to 0 as n + O. Also, most works Detossez et al. 2020
Guo et al. 2025 liduka 2022 Wang et al., 2024a Zhang et al. 2024b Hong and Lin 2023 only establish an upper bound on the average of gradient norms over the time of iteration. In contrast, we directly bound the loss term of the last step to (1). Going beyond convex loss functions, we establish the bound on 4-PL functions, and we focus the constant learning rate schedule.
30

---

E
Constructing the Working Zones
Note that it is generally hard ensure some properties that are crucial the feasibility of our analysis, such as the J-PL condition or the well-definedness of preconditioned gradient projections. However, this becomes possible when we constrain the discussion inside some local neighborhood of manifold. So in this subsection, we construct "working zones" around any local minimizer manifold I such that iterations inside the working zones will be captured by the manifold and obtain certain properties that support the analysis of slow SDE.
For any r C IR being nonempty and R let || denote the Fuclidean norm. We denote the distance from 0 to as dist((,T) == inf 0 Cl Note that when I is closed, the infimum attained (i.c., there exists Ç* with 0 1= = dist(0,1)).
Definition E.1 (Neighborhood of Manifold). For any manifold and positive constant c, the e-neighborhood of I, denoted by I is defined as the set of points such that
dist(0,I)
Definition E.2 (Preconditioned gradient flow, restatement of Definition 4.1). For any differentiable function and any matrix S E Rdxd, the S-preconditioned gradient flow of L is the ordinary differential equation
d0(t) -VC(((()())
dt
When the objective C (or the time dependence of 0) is clear from context, we may omit it in the notation and simply refer to the system as the S-preconditioned gradient flow or the gradient flow preconditioned by S.
Lemma E.1. Let C1. > with C1 C2. and let C R + R be p-smooth and satisfy the μ-PL condition. For any symmetric matrix S with C,I 3 3 C2I. consider the S S-preconditioned gradient flow of C starting ((0) 0g. Then for any T 0,
202
( - 0ol12 s
v [(00) L-
Civ2p
where C infg L()).
Proof. Since C₁I = S 3 CzI, have |SVC())] I C2 (2L() and (VC(O) Z G [[VC()))1 for any 0, which implies
Ci C2
VL())
Then plugging in the above equation gives that, for any T
dt ((O(1)) L- (L(O())) - VL(((()) d0(1)
d
dt
C1
dd(t)
(C(O())) - |VC((())
202
di
Ci (C(O(L)) -T) 22 L*)
a0(1)
2C2
dt
22
d0((t)
202
dt
Integrating both sides gives us
V2nC
T
d0(()
L()))
-
20₂
dt
V2nC
100
-0(T)l2
202
The above equations complete the proof.
31

---

To avoid ambiguity. all comparisons between vectors and scalars are interpreted componentwise. Specifically, for v E R2 and a scalar с we write v - c (resp. v < c) iff i S (resp. v; < c) for every coordinate i. Equivalently
2 e (resp. e > means 2 c (resp. vq c) for all i. In particular the notation S S c means 0 < 04 S c for every i, which is the convention used in the sequel. Recall that IR22 means {e
IR. 2 0}. In the sequel, we slightly abuse this notation
such that for any subset R. By means {0
RD, vs E I for all [D]].
Lemma E.2. There exist constants R, R2 > 0 such that for all k 2 0, 0 S S(v) R2I almost surely. Moreover, Lipschitz on IB [o.
UR S
R and
all noisy gradients (0s( are uniformly bounded by constant R. is also bounded. Combining with the condition that Vggg T) 2 0 for all 9, we have V (0,) (0x) E RD [.. for some constant Rq. Since Ua is an exponential moving average previous V (V, for all & 2 0.
Proof. From Assumption 3.3
Hence
is convex, we have uy E RDD R,. terms and R (0. R.)
From Assumption 4.1 S is p.-smooth, hence both S and VS are continuous, thus are bounded on the compact set RD R.. The boundedness of 5 gives the existence of R2, while the boundedness of
VS gives the Lipschitzness of S.
We continue to the notations Rí and R₂ throughout the following part of the paper. Note that for all optimizers listed in Table I setting R₁ = R² is sufficient. Another thing to clarify is the relationship between the Rg here and the stabilizing constant used by optimizers in Table We will call it Coptim here, so as to distinguish from the c notations that represent a distance (for instance, the in Definition 1. or Lemma E.3).
Remark E.1 (Relationship between Ro and Coptim). Setting R2 := 1/ 1/Eoptim here is theoretically enough for the requirement in Lemma E.2 to hold, but will introduce a large constant the proof since coptim is very small in practice; However, in practice the gradient noise is also very likely to keep v away from zero, thus the operational Rg that governs empirical convergence usually much smaller than the worst-case 1/ C/Coptim.
Now we are ready construct working zones in which nice properties are ensured to benefit our analysis. For all e > 0, define X := _ x D R.) as the set of AGM states (0, v) where 0 lies in 1e
and 0 < S R1.
We construct nested working zones (T Tes) in the following way:
Lemma E.3 (Working Zone Lemma). We denote the minimal distance of I and any other local minimizer manifold as eq. There exist positive constants E1, E2, E3 such that e < 62 E3 < Eg and r9, s satisfy the following properties:
1. L is p-PL. in t for some constant μ 0.
2. For all matrices S E Radd with 11 I = S 3 L1 starting from any initial point Oo 1 the gradient flow preconditioned by S converges to point in r.
3. Under Assumption 3.1 and Assumption 4.1 the function F x - -R. RR, (0, + $ (0) is
C³ on X
Proof. By Lemma H.3 in Lyu et al. [2022], there exists an ng-neigboorhood of I' where L is p-PL for some μ > O. WLOG let E3 E4.
We prove the second property by contradiction. Let C2 1/Rg and C2- = 1/e. Let e2 be some constant such that E2 + V Ge2 < E3. For any starting point Og E 2 and any preconditioning matrix S satisfying C₁ I 2S S C2I, assume on the contrary that the preconditioned gradient flow starting from (0) = 0, will leave Is at some finite time. Then let T inf (t: @(t) & r cc.
Using Lemma E. I and combining the J-PL condition, we conclude that
|0, - ((T) <
V2
2Cs
V C(Oo) L. <
V2pC1
2C₂
V
p 00
00112
=
P C2 1
C
10
001
32

---

for any 0t T.. Hence ((T) E r which is a contradition.
Next begin the construction of T with Assumptions 3.1 and 4.1 Define function f(0, v) R&+D + R2+D
f(0, v) = (-S(v)\((()))
then f is Ca Rd x RD [0,Ri]] Let i be a constant such that > 62- Substituting fo
V F2 + d R1. To (On. vo) such that each entry of Uo is / and 0g be arbitrary point in I, and B X2 into Lemma B.4 in Duistermaat and Kolk [2012), we conclude that there exists some constant 5 such that the mapping %a 0 v) defined by: d@(t)) = - (())))((()), 75 v) = (())
(0)= = 0,
df
is well-defined and C+ on X Note that require slight modification of the original proof since B is now a factorization of a ball and a hypercube instead of ball, but the convexity of B is preserved, hence the modification is trivial.
Note that the constant can be independent with 0g to fulfill the requirements of Lemma B.4 in Duistermaat and Kolk (2112) since || /V and || can be uniformly bounded. Take é1 = 0.9e2. then for any 1 a small open neighborbood of 0 stays in the eg eg-ncighborhoods of two different points on I. Taking union of all Oo E I, conclude that is Ca on X9, Finally, we use Theorem 6.4 Falconer 1983 to conclude that F(O, v) == & s (0) is Ca on xO
F Proof of the Convergence of AGMs
In this section, we aim to prove Theorem D.2 and the first part of Theorem D..1 Specifically, for some constant = 1 O()), we prove that the loss value of AGM converges to Ö(, + m)
within K steps with high probability. If we substitute K = (1 ( log + this will recover the first part of Theorem D.I However, this convergence analysis works for any K = (ppoly(1/n) and substituting K On will give a high probability guarantee that the iteration stays near manifold in the whole scope of our analysis, which helps the proof the second part too.
First, we introduce some additional notations that will be used in our proof. In the AGM framework, an algorithm starts from initial state 0,. and we set my k 2 O, we use step l: + 1 to refer to the process of obtaining the noisy gradient (0x(0)) and then m 1 and 0g+1-
Rd, vo E RD For every
For any h 2 0, to simplify the notation, we denote that
gl := (Ox),,) 2k = (x - C(0s) Z(0x), U == S(v))g; u S(v)) mm+ Φx = Ds (0x).
S, = :===S(vu).
We use time k to refer to the time right before step 1: + happens, i.c. the time right after we get 0 We also define (Fx) as the natural filteration generated by the history of optimization, where each FR (fo. 2o,- Zh-1) 1) can be interpreted as "all the information available up to time k We use the notation E& to denote the expectation conditioned on Ft-
To start with, we prove that the descent direction of each step does not veer off the direction of preconditioned gradient descent, and the mismatch term can also be constrained by list of martingales. After that, we can ensure a decay in the loss function every step, with some small perturbations that can be dealt with using Azuma- nauma-Hoffdding's 's inequality.
From Lemma F. throughout LeaigghhhttLmmmmaaa Lemma F.3] we will assume that the loss function L satisfies J-PL condition at each iteration step, which is automatically satisfied in the setting Theoee D.2) follows directly from the result. After that, we argue that if the loss function satisfies Ji-P.. only within some local neighborhood, an AGM starting near enough to the manifold will stick to the manifold with high probability, which leads to the first part of Theorem D.
Lemma F.I. Let L satisfy Assumption 3.2) Define in == 82 (1 32) 1 (V (gk-19T There exist a constant C1e and a constant C1a independent of L, such that for any k 2 1,
(V (0x-1), Uk) VL (0x-a) ((6x) (0k-1) Yk X
where Ys and Xx are two Femesaaaaale random variables such that:
33

---

1. [& S Cia V2(0,-1)]2 172 a.s.
2. [Xx] & Cb V vL(00-1)ll a.s., and Ex [Xa]] = O.
Proof. We first peel the S())) part off the S(vk) term:
(VL(O=-1),U = (V L (0x- ((0)
= (VL 0 8- (), (V0
S((x)) gk-1)-
Define Y3 Yx (VC (0x-1), ((t
then it holds almost surely that Yul & [V (()()))) S(tx)) 9x- g2-1
Since is Lipscitz, is linear and
00 0kll2 (1 - 32) Ek-1 V (9k-19 - (9k-19h-11112
we conclude that Y& & Cia (V 17² a.s. for some constant Cla- The rest term (VL (0x-1), S(ox)) (( 1) can also be decomposed into deterministic part and a random part as:
(VL (( (())) (VC( (VC(0)))) +
(2)-1, (Ox-))1) V (2)-,, T VL (0x-)))
Now only need let Xx ( T VL (0x-1)) It's easy that E [Xx] 0 and
[ S C1& (i a.s. for some constant Cib. Finally, note that CI is the multiplication of constant bounding the magnitude = and R2 which bounds ||S||,, which is independent of C. This completes the proof.
Lemma F22 (Descent Lemma of the AGM Framework), Let L satisfy Assumption 3.2 holds that
For any k 2 1
- -(0x)-((0)-1) C2n² n1(1 - -B1) Σ
3 ak-i (VC(0--)))
1=1
for some constant Cz.
Proof. From the smoothness of L we have
pn) 2 |
((O%)
If k = 1, then my (1 31 11 So Uy = (1 BI) and the statement trivially holds as long as
C 1 1 If l: > then the (VL(0,-1), 24) term can be expanded
(VL(On ux) (V
= S((0)))),,((v))
(1 (()
31 (V
= =-3 31 ( 31) (VC((x-1)),
81
-
31
5(%-1) m
S - (VC(0)=)))3(0nn))))
(1 B1) (VC(O=-1), Uv)
+ 31 (V VL (
+ 31 VV (()((=))) S(v
1-
Note that a single step of update on 0 and v is small since
0: 0z-1 = 1/2k
U4 Uk- = -2 + (1 32)V (gk-19T_,) [T Uk-1
(1 - B2) (V (9k-19T 1)-1),
34

---

which implies that | 0x-111 = O(n) and 11 - = - We then leverage the
smoothness of VL and 5 to conclude that there exists some constant C2 such that
-(V2(0x=1), u - - (VL(Ox=2). uk-1) - (1- (V 3 C2n. Giving that Ug 0, we can expand this formula iteratively
(VL(0x=1) ux) - - (VL(0 (x=00=2), -1 (1- 81) (VC(O)-1). ( + &1 C2n
s -
(1- B1) UVC( B/02n
$ ...
$ - - 31) Σ
8 (VL(0--)).,, B1 1 E2n
i=1
B1
$ 1 - 91
C2n (1 (1- (1) Σ gi-t (VC(0,-)).U)
Plugging in, we get
((O%) C(0p-))
31
O2n² +
PT
2
| 12 11(1 - B1) Σ
B (VC(O--1)))))
1- B1
i=1
k
S C2n n1 - -11)
g (VC(O,-)))U))
for some constant Cz. Lemma F.3. Let L satisfy Assumption 3.2 and assume that p-PL condition satisfied at all 0g
where & 2 0. Define := 2np[1 PL) For any l 2 0, we have
Ho
((0s) L' S (C(00) L*) + n11 31)
X;
x-33)- 3/-1 +
k-j
Czn
1=1
i=
for some constant Cj.
Proof. We start from Lemma F2 and plug in Lemma F.1
((0%) L(0k-1) S C212 1/(1 -31) Σ gk- (VL(0)-1),U))
= C221 - nn(1 - -1)) Σ 8 ak-i (VL(O,)) - -x).
1=1
Since |V]] S C1a VL((0-1) 1/2 for every i, the effect of Y negligible:
n11 B1) Σ s 'y: I Chans
max (I]V(()-))))))) = o(1²),
==0
and we can absorb it into the Cy7j² term to write out that
[()x) _(On-)(0)))))) S Ezn -n - Bi) Σ - ((0(0=-))"SS,)VV((0) -X.)
(111
for some constant Čz. Note that S(6.) - Hn for any i, hence
so (L( (( 1/ IVC (0:-)11)
* (()x) L(0 S Cz7z² - (11 31) Σ
g (1
|VC (V X.
(4)
Ro
35

---

Combining with the J-PPL property (V,(01-1) 2 24 (C((0-1)
- - we have
((0%) - C S &On -(0 - -C- 2mp111 (1)
Σ 8 (C(0-))))
Ro
+ 11(1 - B1) Σ Bf Bl-iX,
im1
I Czgs² (1- 27 - -B1)
k
(C(0R-1) L*) + - (1)
Σ
8
Ro
i=1
= C + (C(0 (²) + 31) Σ 8 X
Note that we can expand the [(0 -- term iteratively to obtain a generic formula for [(0+)- -C"
k Σ
f(O.) L > (2(0xx L*) + 171 -1)
Bk-
X; + Ezup
1=1
(C(00) -C") 11 - 81) Σ
k-j
8 - X, Σ
j=1
i=1
j=1
S 5* (C(00) -*) + (1-1 - 31)
Σ
X;
k-j
1 0 + Csn,
1=1
j=
where C3 = C3
Ra
Corollary F.1. Let L satisfy Assumption 3.2
There exists a constant Ö independent of L, such that
Vk if
|V(00-1)122>; C,
then with sufficiently small r), we have
<(0))<
< L(0x-1).
Proof. Note that Equation (4) can be obtained without PL condition:
((Ox) L(0x-1) < Eayz² - nn(1 (11) Σ s [V x,
Ro
From Lemma F.l |X,| S Cib [C(00-1)]1 Vi S & where Ci is constant independent of C. So
1
RoC
Ro [V - X, 2 R
VC Cs [V 2
4
where the last inequality uses a² 2ab 2
1 with a = [ / and = Cibv Ro/2
Let - = [VL(0 then
b
1 G - -CNGG--1
C
k-1
Σ
8t i
G.-1 - X1
-X1
2 x1)
16
gk-i
Ro
Ro
4
i=1
R C2
B1
2 Rg
G
C10Gk-1
16
4
1
B
As long as 1 G C10Gb-1
C.
81 I-pi
> 0, small can ensure the loss strictly decreases
Ro at this step k. Set
C RoCIo
1
1+
Vi
then any Gt-1 δ meets this requirement. Moreover, C. depends only on (Ro, Ciby B1) and is independent of C. This proves the corollary.
D
36

---

Lemma F.4. Let L satisfy Assumption 3.2 and assume that p-PL condition is satisfied at all 0% where & 2 O. Let K O(poly(1/m) and ler
w ((0, 1. kk 1} (0, 1)) - R
be a function. Let (Xi) be any martingale difference sequence such that:
1. X; is Fe-measurable and E [X,) 0;
2. [X,] C1 (V[ a.s.
for any i [k]. for any [k] and E (), it holds with probability - 8 that
L(0:-1) L S (1(
then vo E (0, 1), with probability
8, we have L(0,-1) L- S (t, 2x) for all E [k]] and that
ô 4
Xi
S C+
Σ
28 2i2l
i, 2k
log
A=
for some constant C4.
Remark F.1. The (Xi} here may not necessarily equal the {X,} desined in Lemma F. we just make it general benefit future steps. fact, when we leverage this lemma later, we will multiply that of Lemma F.1 by some scalar E (0, 1).
Proof. Note that Σx_ - is a sum of martingale differences. Moreover, since L is p-smooth and 3CIs S.t. every [x] is bounded by C1 VV 1)l12 (Lemma F.I. we have
|x S Ct (VC(0,1))
L*)
S 2Ciss/2pui,, 8') if ((0)-1)--2 (ii,
Since C(0,-1)- S (i, holds with probability 1- 8' instead of probability 1, we create a new martingale difference sequence that masks out all the positions that exceed the bound. Specifically, we define X as:
X
{or
[x
if ((0,,)) - L- S ('),
else.
This ensures that X S Cib / (p 8²) a.s. Then Azuma-Hoeffdinggs inequality gives us that for any e',
x 1 2 6
<
2exp
4
G2 2k 21
22 pu
denoting the right hand side as & gives that for any 8, with probability 1 2
4 z. 16 2k pi) 8¹) log
Let = by union bound, L(0:-)) L < + (i, 전Φ for all [k] with probability 1 1- which also implies X X2 for all i [K] So with probabilty 1 - 8, the following two statements hold simultaneously for all i [k]:
C(O,-))
2k
37

---

and
k
k
X,
70>22 pulle 87) log
1=1
5 2k
C₁
2i
zi
i,
log
where C4 = 201byp.
Lemma F.5 (Convergence Bound of the AGM Framework). Let L satisfy Assumption 3.2 and assume that J-PLL condition is satisfied all 02 where k O. Let rj be small learning rate satisfying = 31/(1 2mp(1-21l) & 0.95. Let K O(poly(1/n)), Under mild restrictions on K, for any k < K, E (0, 1), it holds with probability at least 5 that
Ro
K
((Ox)= - C S Csa -7 (C(0)) L*) + Cs 77 70g 5
for some constants Csa. Csb-
Proof. Denote Do == C(0g) -", and denote the bound with 1 8 probability (k(k 0) n*t Csa Do + Cany log 3, where the constants Csas Csb will be specified by later. We prove by induction. When k: = 0, the inequality
K 2 Do
Csa + Cssn log
holds trivially as long as Csa 2 Now assume that the statement holds for 0, 0, k 1. From Lemma F33 we have
k X;
((0s) - S (C(00) - (1
1 B1)
Th k- B + +Cn.
1=1
j=i
We can bound the coefficients by
8
Σ T
= k- Σ
S San=i k-i
1
1 81
20 k-
where the last inequality is due the assumption - < 0.95 in the statement. Let X
& is also a martingale difference sequence and 20- X,, then { { 1=1
X, |x,| S Co /L(())) a.s.
From Lemma F.4 with probability
6. C(0,-1) - L < vb
holds for all E [k] and also
2k
holds that
4
X,
Σ
2i
log
2k
1=1
38

---

The above arguments give
A
11 - 31)
Σ
181 88-
k
I 220Can(1-B3 B1)
Σ
7
2k 2igh
i. 2k
log
4
4 Σ
2kK
20Can B1)
222-21 21
DCSs Do + C367) log
log
6
k
2K2
4
S 20Can(1 - 81)
Σ 22k-i Do + Σ
2k 222-22 2. log
V log
8
1=1
* CooDD
Cson
2K2
S 20C4n(1 - B
+ 1 1- 72
log
log
1-y
5
>>CoaDo 1-
Cab
2
4
- 20CAn(( B1)
(v
+1
+
log
log
1 1-
8
As long K max 4) (which is mild restriction on K), we have log 2k² and log S log Plugging in 1-7
log $ s 3 3log² K
Rg
7. we have
2
111 - B1)
X;
i=1
j=1
(Vauc(
Csa Ro
3Cst Ro
200Ca - 81)
| M7 Do log
+ mV
log²
B1) -B1)
8
2μ(11 - B11)
Ro
K Vin log
60 60ss B1)
K
1
+ 10C
7/ log
1
11
K
SOoo Do + Coan log
8
where
Cse Coa = 5C 20 Ro 31)/u + 10C4 RCro B1)/u.
B1)/M,
Now as long as K 2 eδ (so that log 2 1), we have
S ((0k) lr S² (L(0g) L*) + n11 - 31)
X,
Σ.
7
x5gg)/
+ +Cs
1=1
j=1
K
S 5* Do + Osc Do + Csan log 8
+
Csn
K
s (Osc + 1) Do + (Osa + C3) 1/ log
8
To complete the induction, we need Csa: Css satisfy
{
Csa 2 Csn + = =5C4/
2Con Do Ro(1 91)
+ 1,
Cssb
2 Csa + Cs 5Cs
205n Do 91) + 100 /C RoCs + C3.
Notice that the right-hand side grows at the rate the square root of Csa and Csb. so there must exist some feasible constants Csa and Csp.. Summarizing, under mild restrictions K > max (202, ed, 4).
39

---

the statement [(Ox) L S 5haCsa + Cosnlog k holds with probability 1 - 6, completing the induction.
Remark F.2. The assumption in the statement, 8 S 0.95, is very mild since 31 S 0.9 Assump- sion 4.2) and 1 - equals constant multiple of n. with small dj this condition is very easy to satisfy. Moreover, similar to Remark 4.1 the assumed threshold 0.95 can be replaced by any constant below 1, and the order of the convergence rate will remain unaffected.
F.1 Proof of Convergence-Relatee Conclusions in Appendix D
Proof of Theorem D.2 This direct corollary following from Lemma F.5 The loss function is
global J-PL in this case. Setring k K, letting >k O(1) gives K O log 7) completing
the proof.
Now we move on prove the first part Theorem D.I] The main difficulty comes from the fact that C is only guaranteed to satisfy μ-PL condition within some neighborhood p The iteration, once getting out of that neighborhood, cannot be characterized. Hence we need to bound the probability of that event.
The trick here is construct proxy loss function 2 that, agrees with L near I' but has "wall of quadratic functions" upon L further away. L thus satisfies (μ, L)-PL which allows us to use Lemma F.5 If the losses at all steps are small, this ensures that the iteration never leaves I where andd C are identical. We formalize this idea the sequel.
Lemma F.6 (Tubular Neighborhood Theorem; Theorem 6.24 in Lee 2012)). Ler C Bd satisfy Assumption 3.4 3.4(in particular, I is a t compact embedded submanifold) Then there exists TT > such that, writing NT for the normal bundle,
V = u) E NT ||ul22 < <r);
F:V R.d. Fly, u) = p
the map is a diffeomorphism onto the open tube U == l Consequently, the nearest point projection P U + is well- and C and every 1 E U can be written uniquely as
0 P(0) + v(), v(0) Nr
Corollary F22 (Smooth distance and unit normal the tube). In the setting of Lemma F.6 let U= F and write E (0) = (P(0), (0)) on U. Define
n(0) v
r()) = dist r) (1)(0) 12,
(0 (UE
|2(0)]12
Then and are - U \ I, and
Vr(0) n0),
VO U..
(5)
Proof. By Lemma F.6 P and are - on U, hence and n U \ r. Set g r(0)2 [[() By the chain rule.
(g = 2 (Vv(O)) v()).
From the identity 0 P(0) (()) have
Vv(0) I - VP(0).
Since P(0) maps into T and (()) NP it follows that
(VP(0)) T v(0) O.
hence
(0( = 2(1 2(1-VP(0)) v(0) 2v(0).
Therefore, for 0 U (so r(0) > 0),
Vr(0) = 2 2r(0) (g(0) = v(0)
= n()). |v(0
which is Equation (5)
40

---

Lemma F.7 (Nonobtuse angle between VL and the outward normal), Ler I satisfy Assumption 3.4 and let L satisfy Assumption 3.1 Assmmpion 3.2) Then there exists a constant (0, Tri such that, for all 0 E I with hearest-poin projection P(O), distance r(0) == ||0 P(0)]l2. and outward unit normal n0)) = (0- (())))r(0), we have
(VL(O), n(0)) 2 20.
In other words, AV n() S x/2 on 1 AT.
Remark F.3. This lemma also implies that + ((P(0) tn())) is non-decreasing on (0, (0, since d [ tm) (VL(d) tn), n) (0, r). To put it vividly, I is valley, with I being the floor at the center of it.
Proof. By Assumption 3.4 each Ç E is a local minimizer C, hence VL()) = 0, and V² is positive definite NçI, with uniform lower-bound >> 0 of its eigenvalues N&T (from the compactness of r). Let 7T be as Lemma F.6 For r write & := P(0). : l|e Φ 2, and = (0- / N&T. Since L is C -mosohhh,
perform third order Taylor expansion:
V 0 V + V2L()) - Φ) + +
||0 & (
where ((3) is constant independent of 0. Taking the inner product with n and using VL(Φ) and 0 Φ rm, we have
(VL(0), rn V2L(()) + (0(3), n) 2 mr C(3,
Choose == min(rr, m/C(s)), then for all S we obtain Equation (6)
Proof of the first part of Theorem D.I First we construct tubular neighborhood around T and introduce some notations. By Assumption 3.4,
the unique set of minimizers of L in some neighborhood U of r. By Lemma F.7 there is E (0, mr for which the nonobtuse condition (VL( (0) 2 holds on I By Lemma E.3 there exists E3 such that p-PL on I'ss Shrinking é3 if necessary, assume é3 and
53 C.
Throughout this proof work inside the tube n given by Lemma F6 In particular, for every E 7T we have the nearest-point projection P(0) E I, the normal offset v(0) E NN the distance and unit normal
v(0)
(0)) == dist r) = [|x(0)]22, n2()) = [((()
and (on U \ () the identity Vr(0) =0) from Corollary F.2
Next, recall the constant 6 constructed in Lemma E.3 such that 0 < 61 E3. Define the "gap level"
Lm := min (inf { [(0): r()) [e1, [aleea]]]11 L
The set (0: r(0)=
{g]} is compact and disjoint from r, hence f. > L
Then define the proxy objective É: Rd R by
[((),
dist(0,1') S E1,
((0)
C(0) + (O.st(0,1) - =1)² dist(0, > E1,
where a large constant satisfying C 2 p. Note that for E Rd.
[(0) < Lm = dist(O) <
so on the sublevel set (É )m) we have Ż Ž=6 L.
Now define
C (es - =1)² 2 +.,
L =
2
41

---

and we prove the core property of the proxy loss function: is (μ. L)-PL. First we consider the case
r On 3 the distance function r(0) dist 0, H) defined. Using Vr(0) n0 from
Corollary F.2 and the nonobtuse condition from
Lemma F.7 for r(0) >
VŁ(0)= V2(0) + C(((0) -1) n()),
VC(0)] [V2(0)]]+2 -1) (VL(O), (VV())))))))) (r( e1)² =1)
z [V
=)) -
Since is p-PL on r
C 2 (r((0) -1) > 2 minff (E(())))))()))))))))))))) -*).
|V(()) 2 2p( - --)+ 22
For r(0) 1 61, -= L and the p- inequality holds trivially. Our choice of yields
|V 2 2u(((0) - for all O E gr,
Lr)
ie., is p-PL on I9.
Combining with the fact that ((0) L i & T we conclude that Ł is (1e, L.-PL.
Finally we come back to prove the main conclusion. Let C the constant constructed in Corollary F.1 Note that for 0 R \ re,
|V(())
+V((0)+((r(0) e1) n(0)ll2 CC <Č 55
- +[[(((()) 9) (0)ll2 S C (from Lemma F.7
- r() < c te..
Increasing C if necessary, can let C/C (e3 =1)/2, which implies
63
cm C
Take some Co L such that Lo Co-C" (Lm Take constant < E1 such that all loss values inside 1 are bounded by Co- There exists sufficiently small learning rate 7 such that (let K [T 1)m- D):
//2/OOsa where Cbe is the constant in Lemma F.5
K 1. Caatt (Lo L²) + Cssn 5 s 0.99 Lr, Vk S K, E
2. 1 R R2 < C3
Cm-
The second property ensures that any single step of update cannot jump from the interior of I to the exterior of e However when 0x-1 T T it follows that || > C. so [(O)) < L(0x-1) from/Corollary F.1 By induction, we conclude that for any 0g E 1o if we launch an AGM from O, and train using C and y), all loss values ((0x), k E [0, K - 1] do not exceed L,
which means the μ-PL condition of is satisfied at all O₂ where k E [0, K 1]. This meets the requirement of Lemma F.5] By Lemma F.5 and the first property of n, further conclude that all loss values [((k) do not exceed 0.99Cm. Finally, by noting that
((0) < Lm + 0 DEI and Ž(() L()),
we conclude from Lemma F.5 that: For any Og E [ if we launch an AGM from Og and train using L and 77 for any k < K, d (7,200, 1). it holds almost surely that 02 1 and it holds with probability least 1-6 that
K ((Ox) L Csa * (C(0o) - L*) + Css n nlog 8"
and takes K2 = ( - log ))) time to reach L(Kg) - -L" 0 (1l08 1/)
completing the proof.
42

---

G Proof of the SDE Approximation of AGMs
In this section, we present detailed derivation of our slow SDE approximation of the AGM framework as formally stated in [heorem D.I In Appendix G.T we give some lemmas about the properties of the adaptive projection 4,s, introduced and explained in [Section 4 Then in Appendix G.2] we show that the iterations after convergence would continue staying near the manifold. Following that, we further calculate the moment change of the projected parameter 4 (.,()))) during a "giant step" in inppppediix G.3] which allows us to conduct an weak approximation in Appppdix G.4 After going through above paradigm, we finally prove an equivalent Theorem G.2 the SDE approximation part of Theorem D.I
Our slow SDE starts point after the time of convergence, for simplicity, we will "shift the timeline" in this section.
Remark G.1 (Time Shift). To simplify the notations, we redefine Og and vo as follows. Starting from Appendix G.] Oo and vo will no longer represent the parameters that are initialized at the actual beginning of training. Instead, they represent the 0g, and UK. yielded by the first part of Our SDE approximation then describes
Theorem D.I (01, U1) denoting (0+++1, UKg+1), and the dynamics of AGMs after reaching the state vo)-
Remark G.2. Recall for any time step k that Φb == s()))) With the "time shift" described in Remark G.T the time steps before Ko will become negative. However in some parts of the following calculation, deal with the first-order momentum we still need up to Ollog ") past timesteps. Without loss of generality, we assume that at time Ko the iteration has already converged for time Olog )) ie, vKo Olog b & k S Ko.
17 log 1
C(0x) L- =
(V
00 -bx122 O
IVVC(0)) O
(V 1 1)
If not, we simply increase Ko by (log 1)) and the argument in the proof of the first part of The orem D. still holds. After the time shift, the range of k above will become -O(log 1, < k < O.
Therejore, negative timesteps may appear in the derivation below and they are not typos; We just need their three properties above to control the order of some terms.
G.1 Lemmas for Adaptive Manifold Projection
Before we characterize the projections, we introduce some properties of the preconditioned projection function in this part.
Lemma G.1 (Adaption of Lemma C.2 in Li et al. 2021b)). For any E Rdd, and any p.d matrix S E dRdxd d it holds that 00 (z) = O, and
8 ( SVL -
Proof. We consider trajectory starting from s= ac, with ODE dz(t) by the definition of 4ss we have 52(x)) = 4 (t)) then we have
SV ((x((), thus
( -00 (
dt
Further, we take the second derivative of (x)()) with repsect to
s (=))(((((()))))))))))))))))))))) (VL (oa()))
di2
Taking 0 completes the proof.
Lemma G.2. For any E I, and a p.d matrix S, the following two identities hold:
43

---

For all E (T),,, (
For all T (T), s Su 0.
Proof. We first prove the first identity. For any v E (T))), let {v((), smooth curve on I' such that v(0) a and de(t)
( 2 0} be parameterized
di === v. Since v(() E T, thus vL(e(()) - 0.
which gives ss(((t)) vtt. Thus we have
du(t) t=
d (
(v)(v(()) de(t)
dt
dt
dt 1==
1=
Plugging
dw(())
de 1== = v gives {z( = v.
For u E T (T) and 2 O, we consider the Taylor expansion of V +V2 at =
VL x + [<))(()
V²((x) (V o(t) = tu + o(t):
where the second equation uses the fact v²(((x) is full-rank on Ta (T). Thus, by the continuity of S&s proved in Lemma E.3 it follows that
8 (x + +V+ SV =
lim 600
By Lemma G.1 0 (a + (()))()))) SVL (V = 0 for all t 2 0. implying that 00 (0)( = 0 for all TT((T)
Lemma G.3. For any x E T, and a p.d matrix S, holds that 00 s SV) = 0.
Proof. From Lemma C.I in Li et al. 2021b), have for (T),), V²L(x)u and for T (I), Lemma G.2 gives that
(e) 0.
The above identity completes the proof.
Lemma G.4. For any x E r, u, E Rd, p.d matrix S, and v E To (T),), it holds that
( 1 = (0)(0)))0)))) (V uo" S - (1
Proof. We define : 51/2, And we do reparameterizzaiio as x' = P-1x, L' L(Px). then we have
00'(4") (P
V²I'(a=) PV')
8 (VL')<z' PPE 8 55 [M] PO
Notice that in the space of æ', the adaptive projection mapping & turns into fixed gradient flow projection. And this allows us to directly apply Lemma C.4 in Li et al. (021b), which gives
03 (2')[, u| -00 (2 V' V2
A slight modification using the above transformations gives
(0 Pu] = ( (
( Sa Pu).
We now redefine u Pu, v Pv. and we organize the above equation
20 - 1
We completes the proof.
44

---

G.2 Iteration Stays Near Manifold
Now we begin the final preparations before deriving the slow SDE near the manifold. Note that
in the end of convergence analysis, the total steps equal K = 1(T + 3 and the converging step Kg O 1 log
n. ). So after time shifting, the high probability convergence of L(T + 1) 2
O log +) 7 number of steps in our analysis. Let / be some constant in (0, 0.5), whose exact value will be specified later. First, we bound the movement of projected steps by showing that Φ shifts no more than (170.500.59) within AK == [n-1- steps, demonstrating the "slowness" of the dynamics of AGMs after the projection.
[T steps are ensured in Lemma F.5 Now denote K :- [ be the total
Lemma G.5. For any d (())ly(n)), with probubility 1 1 for k E [0, (0,K-AK) Ak E [AK].
|1 2el12 0..005-0.50 log
1
np
for some constant Co.
Proof. Recall from the proof of the first part of Theorem D.I that O₂ stays inside 1st For any E [0, K] almost surely, and from Lemma E.3 that s is C4 on X == r x [. Since x is compact, (0) is then bounded and Lipschitz on A Similarly, 00 s(v) (0) is bounded and Lipschitz on X For any l E (0, let k = k 2 log,, 7, we have:
Φx+1 Φx s(( (0x+1) sv. (0x)
17² 1/ 1
(vvg) (0x+1) {s(oz) (0x) o
(0s) 0s) + 0
v7 log
1
9
1 7)
O (1)
+ 0
17² log
50 (0)
+
n" log
1
where the second equality comes from the fact that one step of update on of O(n²² and the Lipschitzness of S and Φ, the third equality comes from ||0++1 - 0x O()), and the last equality follows from the boundedness and Lipschitzness of Od. We decompose my as:
m (1 (1-33) (1)
3 (VL(()) + =z)+0017)))))
3 (V[(0))+ + (nlog 1)) + (1 B1) Σ 2 + - 0112)
(1 31)
A key observation is that 0 S(r)) (()))(e)) v ((0g) 0 from Lemma G.3 which allows us view Φk+1 Φi as Σ žt.: + (17 log 1,) where εạ. = s(ve) (n)) 31 ) 22) Note that ž is F and its mean is 0, since E..: just applies a linear tensor transformation to z;. If we define constant Co. == sup s()))s())
(v, 0) E X } B1) R2 that is independent of and i, then |ا is almost surely bounded by n 'Can ||
For any K E [0, K - AK] and Ak E [A]] we have
k+Ak-1
ΦE Φb =
Σ
( - Φ,))
J=k
-+
z O (1fl0g 1
Σ
J k
imj- 21ggo 7
45

---

(RggH H + + logs, n) log 2 g minfk+ 10--Ak1110 A1k-111+1o8,, fi}
b+4A-1
Σ
2/ 2
2= log, 7
j=
min k+ Ak 1. logg, "}
Denote Zi :
أزر
2, then each Z is linear transformation of so it is
with zero mean, and also |2:1l22 S 17 then gives that for any 8 = =(poly(1)))
Cca 1-8
Coa a.s. Azuma-Hoefiding's inequality ||: 11 S 17 1-8 8
with probability 1 ó,
Cea
ΦR+AK Φx S /22
I <Css V 17 11 log
for some constant Cas- Finally, plugging 8 [0, K AK] and Ak [AK] gives the theorem.
K-AK and taking union bound over all k
E D
With the concentration bounds so far, we can show that the dynamics behaves "well" during the whole iteration, and formalize this idea below.
Definition G.1 (f-good). For any (poly(nn) and any step K [K], we define step K to be d- 5-good if and only if the simultaneous establishment of the following statements:
1. For any ls E [0, KJ. db E and |00 - -0lll2 S Csa / log np
2. For any k E 10, - AK), Ak E [AK), |1 Øxll S Cssn)..-0033 V /log n
10,K-ANT
Here Osa and Css. = Cov2 are two constants.
Lemma G.6. When 1j is sufficiently small, with probability 1 17,00, the event -ooog holds for any step K in [K].
Proof. Denote = 1) 100 From Lemma F.5 I -
with probability 8/2, all E [0, K] satisfy
(00)--++CCG log 2K 8 Note that Du := -(0,)-C* is of On / since time 0
now refer to the time after convergence. Combining Lemma E. this implies 110, - Φx ll2 S 20,
V2μ01
2K² for any & E [0, KJ. When v) is small enough such that 00 Φb ll S V Osa Do + Css log 5
20, VCaa + Cssy log 2K 3 + njRe 62, any de T with k 2 0 will imply Φ E T,
V2nCi
since 0x+1 cannot escape 1 Giving Φo E and using induction, we conclude that all Φx E T for k 2 0.
When the above holds, the requirement of Lemma G.5 is met. Then with probability - 8/2, for any
A E J0, - AKJ. Ak [AK], we have || Øull2 S -. log 2
176'
Finally, we just take the union Lemma F.5 and Lemma G.5 With log 2K log 3/ < 2log (which mild restrictions since is small), have the theorem.
S 8 log 1 and i
We have proved that iteration will behave well with high probability, but chances still exist that the iteration driven out of working zones and becomes intractable. We define a bellveehaved sequence that manually redirects the iteration when extreme cases happen.
Definition G.2 (Well-behaved Sequence). Denote the event of step k being 11100- -ood as Et- Let
Small be a fixed point on I. Starting from 0 = Oy and ún = Uo. we define a sequence of (0x, 0k m as follows:
m B1m4 + (1 (( =*)
6x+1 = B204 + (1 -3))V(((2( (V(()( - T)
0.+11 : 1s, 0++1 +1,, Baull
46

---

where 1 is the indicator function: 1g = 1 if event & happens and 1g 0 otherwise.
Note that the update of 02 can be written as
Bk+1 02 -7S(0)+1)mm++
- -18 (0x - -)(( + 1g, Paull
where e₁ denotes the redirection under extreme cases. By definition, ex equals zero in the vast majority of cases, and other cases it's still bounded by a constant, so all moments of e₁ are within (r)) which is negligibly small.
G.3 Moment Calculation of AGMs Near Manifold
Additional Notations. To utilize the analysis framework in Gu et al. [023b]. we first introduce some notations needed. Consistent with Gu et al. (2023b), pretend that AGMs proceed with H 1 local steps, as a single worker (without multiple workers). We denote every steps as one round. Next, define "giant step", which encompasses Rgrp 7 rounds, corresponding to Rgpp" H steps. We consider total timescope of f/z steps, which corresponds to 7- giant steps.
T
For any 0 S 8 < Rgrp and 0 S 1 S H, we use 0, (s) and 0% (where k = sH + f) exchangeably to denote the parameter
get the t-th local step of round 8, which is also the k-th global step. Also
note that for any 0 S 8 < Ryp. 0/4 and 0(++1) refer to the same thing. We define the notation o(()
m!
(s) and e(a) in the same way did for O. Furthermore, we define
gr (e) velo) (0()) Sx ((x), S S sb() S(()) -(6)) 266) =
&g (0(-)
a 0, (6)) A = : - 300 Eo = E
(0)), (0)) P1 := 0Φ S ())),, P1 = I Pl:
a = E a 1 = E a (T) B( = 2 (
Corollary G.1. There exist constants Coer Cyb. Cse such that for all S s < Rgp. 0sts H.
Cony
log nj
0
(6-))
Cosg/
1l0g π
(30+)
2(0)
Caaa 0.5-0
log 1)
Proof. Substituting 5 - 1,000 When & holds, this follows directly from the definition of 6-goodness; Otherwise, all & and & equal, and these quantities are equal to O.
Impact of Momentum. Our conclusion regrading to the impact of Momentum on the implicit bias similar the conclusion in Wang al. [2023]: : It does not impact the implicit bias. Further, our analysis is based on moment methods and can give exact error bounds. First, we state some technical lemmas in order to show that introducing momentum will not cause the gradient deviate
much from itself. i.c. [[rr is close to E g.g.). Once this guarantee is established, we can replace m, with g, in the moment calculation simplify it. The general idea of the proof show that if i is close to t, then ((V))(0;-1) will become close to E and if far from f, then the contribution of E[V((0,))) would be negligible in [E(rr.).
Lemma G.7. For any k 2 O, we have
E
VC (
JL (C (0:)
S 2.5
for some constant Cio-
47

---

Proof. We have
vC((k+1) - VC((0)) V2 (((()))00+++ 0x) + 0 -
V²L(0,) 0,) + +0 + O(1lexll))
since || | 0z| 00]1 - I| (0x) m. -exl22 O(n) O(1ex 12). Let k = k logg, (n) be
threshold that is logarithmically close to k, then we have
V ((+)) 0x) ((2 O (10. 0% 01())) 00+1 -bx)
v²c(0; (0i-1 - -Os) + O (17 log (1)) n7) + O(llexll2)
((22((()5 (ix-1) m++1 + O (1p log 7 ) ((1exxl))
Recentering the Hessian term to 0g allows take conditional expectation Eg on (0x+1) m
E V ((0))S (6x+1) m
E V VVC())) Eg [S (0x+1) m
After that, notice that
[ [S (0x+1) met1lll2 = ||Eg IS (Eg [ m 1l1 01 Eg [0x+1lll2)
||S (Egg (6x+1) 1) Eg m 12]12 + 01 Bg (
= O [ (2
O1 - Eg (0++1lll))
=:
since S and in are both bounded by constant scale. We figure out the orders of these two terms respectively:
k
D1 = Eg
3 k k+t mn + (1 - (1)))
ak 91
i=
k
0 (31
logn (1)) (1))
+ Eg
(1 - 31)
gk-
91
1=
O(1)
R3 (1
B1)
Σ
i=k
O(1) 01) (n 0,5
0.5)
since VL is uniformly bounded by 0(1).)) after convergence (see Lemma F.5 And
D2 (1-B2) Σ i=k
B (V(g.).) Eg [V(g.T.)
= 0 (b2 (k -))
1
= o n¹ log
since V is bounded by constant scale. Now combining the above together, we have
- -E v c(O E4 IS (0x+1) ( + 1r log !
+ (EI]eell)))
ODD0 + D2) + 0 (1 11 log - + 0
0)
which concludes the proof.
48

---

With Lemma G.7, we are ready to deduce the closeness between E(mrs] and Elgh).
Lemma G.8. For any ls 2 O, let k = k -2ogg, (7), we have
| 9x+1lll2 S C119 15 log 1 ..
a.s.
Note that this also implies that [| - 2xt1lll2 S C1191 log $
Proof. Expanding Eg(r we have
k
Eg m+++1] = Bí
(1 - B1)
gk- 91
i=1
h-1
k
=(1- - B1)
9: + (1- 31)
Σ
84 3 19:]
1=1
i=
k-1
= (1 - 31)
3
9: + (1 31)
3 VL(O)
1=1
i=
:=E
E2
Note that Eg is neglegible:
R-1
E 1
(1 (1- (1) Σ 3 gi
R-1
=(1 - (1)
3 0(1)
1=1
$ (1 11
0())
i= log (n)
0 (3) loog, (v) -0 (1²²)
and that Eg is close to V(((A)):
k
|| BB2 - [VVL(
(1 B1) Σ 87 -
=
2
i=k
k
(1- 1-331) Σ
E V [v(())-V((.)]
+ +nn²²
i=k
< (1 - B1) (I - k) Cwon 1.5 01??)
(by Lemma G.7
Combining the results of E and Eg gives
EB [m gulll2 - |E 12 + Eg - E[VL((+)]
< (1 - log (n) Cion 1.5 Onn²)
CCup 1.5 log
1
for some constant C11 which completes the proof.
49

---

G.3.1 Moment Calculation Within Giant Step
In this part, we aim to give the change of first and second moments of and é, which is the basis of deriving the SDE for AGMs.
Now there are only a few preparations left before we get into the direct part of the moment calculation.
For all 0 < Rgrp. < 1 v H. Note that | + l2 (1 32) | (gv.%) -b lll
O {2) 01))), so combining with the Lipschitzness of 5 gives
Sk2 Sks |1
o ((k2 - ((k2-k)))) kn)
for any 2 > k and bo k1 (17 ()) Next, we begin our moment calculation analysis, starting from the update in single step.
Lemma G.9. For all k S Rgrp H, it holds that
E
O0++1 E 0 S/Sgg O (1)26-9).
Proof. We write the update rule of AGM under a single step as
0x+1 = 02 - n - EE
= 02 - Sing + Sx (mx+1 - gu) Ek
02 -n n Sogn + (Se= - So 9x
Sa (m++1 - gk)
Rk:
A02
Ad,
where we recall that 6₁ := 11 (0x 175(( (0x+1 ) + 1g, Pmull- We can prove that A0, and Ad₂ are small enough to be negligible in expectation for our following calculation.
Specifically, if l: = then Ad, = 0; and if O, we can decompose E
Ad as:
A = ( ( - S 9k + (St - Sk-1) 91.
= =E ($ - So) VC 0x)] + E [Sx 5x 91
=O((k 1)m² 17 03 + O
O(H Rgp 11 ))²)
15-
Here, the second equality holds since the gradient noise term as step k z4 is conditioned on time k; when Sx-1 has already been determined, thus we can take the conditional expectation.
For 20,, let & = k - 2 log,,, (1), we have
E | E
S (m++1 gk)) O 17
12 log -
S EE [ + (1) (12 log -
E
17 1.5 log
1
10'100" log 1
+
17
n 1.5 log
where the second-to-last equality follows from Lemma G.8
Finally, we have
+ O (5-) + o 1 log 1 17
E 0 0++1
E [
1,$09z
+ O1, 100
E 00 S)S0
+ O (1)
50

---

which concludes the proof.
After getting the update rule of Óz. we then derive the moment change during the single round with H steps. this end, we recall our modification of manifold projection from a "Gradient Flow" manner to a "Preconditioned Flow" manner in Definition E.2
Definition G.3 (Preconditioned Flow Projection). Fix a point 0muit & P. Given a Positive Semi- RR consider the preconditioned flow dæ()) - -VV((x(t)) with x (0) Z. We denote the preconditioned flow projection of x as P (2),,, i.e. (M) lim x(t) if the limit exists and belongs I, and t (2) = fuull otherwise.
Definite matrix M. For
We decompose the preconditioner matrix in the very begining of the giant step as So ((6) PP. where P = Su We then provide the first moment calculation of & in the following lemma. Before that, we first introduce the operator VH.
Definition G.4. Given a Positive Semi-Definie matrix H E Igddd, whose j-th eigenvalue and the corresponding orthonormal eigenvector are denoted by X5 and Vj. We then define the operator VH()) Rdxdd > Ruxd as
vH()= Σ FVAA
d + Xg (10,0,) T.ao.
Intuitively, the above operator projects the matrix into the basis of H and sums up the corre- sponding components with weights x,+ Then we present our moment calculation lemma.
Lemma G.10. The expectation of the change of the manifold projection every round is
Hn²²
E (((+t)) -(0()t))~6(1))
-6(1)
S0 s. (6 (
(PV 0671-5-8)
2
for Ro <8 < Rerp, and
E [(0tt)) - =Ō(n)
for S Ro, where Ro == max
{[
10 log Amaza
1 2 2oggi/3 /3 } and Σo,p P.p
nj
Proof. First, consider the scenario when Ro < < Rgp Let L'(x) =- L(Px), then
VL'(z) = PVL(Px)
T/ PV (PPx)
E((r) PPxPa)
8² (VL')())(M] = PO (PL)
For one-step GD update, we consider an auxiliary process (0))
@ = 0. v)) + (122 0
= 0 -P + 0 0 0.
Similarly, define A, = B[ (+) 2; ( d. (-) Eaa (5) and B((e) Eia ( T) and
(xx) is the gradient flow projection of point æ. We further define o'(s) == (6)((s))
Now we are interested in the update of PỘ', which is
Pθ = P S V..()))))) + 0 (12
One can obviously see the update rule of P resembles the update rule of 0 in Lemma G.9 Now we set O' - 0 then Equation (7) satisfied, and combining Equation (7) and Lemma G.9 gives
(()) (()) ge41 gr+1 - + (n.25-0)
51

---

Notice that the above equation resembles the single update for SGD, which allows us to apply Lemma 1.36 from Gu et al. 2023b) for the update of 0', with loss function (0),, number of workers k = and manifold projection ('))), which gives
PE (
E (B(+tt) (3())
PP30 so (8 PP
(PE,P)
2 (
where the first equation uses the fact that PO( = ( and can be verified with the definitions of Φ¹, 2s. and 0'.
The proof when I Rg is a direct conclusion of Lemma 1.36 in Gu et al. 2023b since the Ru x log our case.
Notice the above equation for the moment of á contains ø'. The next corollary eliminates o' from the formula.
Corollary G.2. The expectation of the change of manifold projection every round is:
( -m(+) -c(t))
- - Ö(1),
Hn S s. ($) + (1))
Ro < < Rgp
E
Ro
Proof. Notice that for the preconditioned projection, also have the corresponding transformation
O (e'(x))))))))))))
Pa 3 Px')P
2 [M] - x ['MP
The above two equations and Lemma I.36 in Gu et al.
2023b complete the proof.
Lemma G.11. The second moment of the change of manifold projection every round is
[ (())) -2)00) {
Hr PIpPP So Eo PI..s So Ö11225-8),
Ro Rge
Ö(n).
S Ro
where Ro == max
[ log /
logi $
} and P. == 0
Proof. According Lemma 1.37 in Gu et al.
2023b we could write the second moment for 0' as
Hr + -()11--8). Ro < < Rgvp
E (6'(6+1) -
(6''6t))
- -(0)))=((
Ö(m).
$ < Ro-
Notice that
E6 0 (600
Pap ( PPOD (b)0))
When Ro S < Rgaps
E ( -
= P (s+1) g1(s 16"t)))))
S EoS I So.
The proof when 8 < Ro is a direct conclusion Lemma 1.37 in Gu et al. 2023b since the Ro x log in our case.
Then give the moment change of o within single giant step.
52

---

Theorem G.I. Given |1000] - |12 = -01 olynlog 2)), for 0 <B < 0.5, the first and second
moments of A Ban) = (bllae)) T6) are as follows:
E| ( 11 = 11 11- 8
S 0ho o 5o (6 " 1S0 + (11)1.5228) O(1),
2
6 RAP)) T|| = 17 1 S
100 ( + 0 + Ö())
where En(0000 ((0) == P S
Proof. First we prove the first moment change as
Bgp
1
E B ( = E
( 1)
-c60)
==0
Ro
RATp
Σ
Bl
(++1)
-Bit +
Blblot1)-
Ro 11
I Sol + Ö111 + Ö(r)
2
The last equation is a direct conclusion of Corollary G.2 And for the second moment, we have
Rxp
Rxp -1
((0+1)
( - (3)(
E
==0
Rap -1
Σ
(686++11 - - (T)
==0
B[( ((()))(((3 +1)
(+(+++1) [
S -
E (6 (0) (0)) So 0(1 1.5560 + O(())
where the last equation uses ([(
(( Ô(17).
Next, we proceed with the updates of v.
Lemma G.12. Given e == 1-02 and we have
x12
E o Rppp)
[ cn -3 (v (2(")) -0.)) + +0 (1655-1.5)
Proof. By the update rule of v. we have
60
(1++1)
(/.0) 0//0 +60)
= =3 00 (0) + (1 (1-32) 32) Σ
Al
-
V
(g(0) g000
(+) -/
H
= -3H 1) (00) + (1- - B2)
Σ
3tt - v ( (,,²0 TT
1=1
Note that
(c) (x)
= E Σ [ ]
E
9 9
E E(6 (1) + x [0)]
E E( 0 033333))))
= 1 + 0 0.5-0.5))
53

---

Combining with the linearity of V, we conclude that
(++1)
(5)
(3 1) 00 + -1-)))) V (2)0)) + +0 (171.500.8)))
E 0
E
0
[++
(5) 3 00 + (1- / (1-B2)
V ( + O (111660050))
To transfer from
60
(0) arbitrary 06%) we simply expand to get the result:
00
(x)]
= 52 00
(0)
+ [11 85 V 1 00) + 0 (1) (.660))))))
] ( + 3 + + ..... 1/6-11111
1,2 (0)
+ [11 / v _ ) ( | 1 3
(1-
82
+ +0 (15555066))) (1
12
(0)
+ (1 /2th ) V ( 0 (1) (5)_05
vo
Thus we have
E 0
Rgp))
0100 cr 1-
cn (v (x(())) - 0000) + 0 (111661.6.))
where the last equation uses the fact that 32 Rum = 1 - 1001- =
0
Also, for the second moment change of é, we get the following lemma Lemma G.13, The second moment change of i over a giant step is
(00 (Rai)) - -000) ] 0
E
-60
Proof.
E[ 11 (++1) - -b(()) (og (((0)) T 1-E] ]
E
(8 1) (1- B2) Σ 3 v V (0.0
3 - 9 0 "))" ))
T
(8 -1) + (1 (3))
Σ
0 ((1 - -3 = (xp?)
RxI
RxIP Σ
|
1bb (Rww) =0/))
(R
(0)
T
(s+1) =b61)) ))
(s+1) ob6")
66
E
0
0
Reep -1
(8b+t))) (s+ - -0(")) (ol -(66)))))))) |
E
s=0
E [1 .b6")) R 00
(s'+1) - =6-")
sps'
0112-89 ))
The last equation uses
El Un - =66²) (o((t1) -b(-)) 1 0(n))
(s+1)
and
( (0( ((++ - -b/
(06
(((''+1)
1 -8((1) 7 ((3-33).
E
E
The above equation completes the proof.
54

---

G.4 Weak Approximation
After we get the first and second moment changes within a giant step, we now utilize the moment calculation to prove the SDE approximation part of Theorem D.I First, we recall our slow SDE for AGMs
dc(t) (c.s()) do(t) c (V(((()) e) dt.
- http://wwwwpptyecccccccccctttttd
We then open the projection mapping (c.s(() as
{ dç= ( stu) st))(()))))))))))))))))))))))))) (C)))) dv(t) = cV u) dt.
(8)
Now it suffices to prove the SDE in Equation (8) tracks the trajectory in AGMs within ( steps in weak approximation sense.
First, we have show that the solution of Equation (8) close in the minimizer manifold Lemma G.14. Let (t) (5(t) T, =T be the solution of Equation (8) with ((0) I, and v(0) R then we have that (()) E T for all 1 2 0.
Proof. According to Filipović 2000 Du and Duan [2006], for a closed manifold A to be viable
the SDE dx (1) = A(X ((), + 6(X (1))d/, where A(-) IR+DD - gid++)x(d++D) and b(-) Rd+D + R are locally Lipchitz, it suffices to show that the following Nagumo type consistency condition holds:
2 Σ (AA 2
p(z) = b(x)
e (T(M), TT(M)
where D[. is the Jacobian operator and (x) denotes the j-th column of A(æ).
Following the argument in Gu et al. 2023b] here also only need show that P (z)s())()))((r))
O, where (z)) = La (z)).
(l,s(a) Σ ([A = PI,s(ar) Σ D 3 ( /
(
= -P > 0 (
S00 ( SE
SE (z. S)]
Notice that, since it is clear from the context, here write = S for short. The last equation uses
Lemma G.4, Agian, applying Lemma gives
P (i,s(z)b(z)) -
P (F) [SE S)].
The above equation completes the proof.
To establish Theorem D.I we give an equivalent theorem, which captures the closeness of X(t) and X, in long horizon. Also, for the proof of Theorem 4.1 it suffices to prove the following lemma, whose proof will be shown in in/pppendix G.5
Theorem G.2. If ( 10 = O(1 log + and ((0) = ((0), v(0) (0), then for a giant
step Rgp = 1, for every test function E C3.
Elg (x (x(nRer))) Elg (x (mm (.76])] -Cn 0.20
1 -yb 1
max 05 G0.rS
where Cg a constant independent of n) but depends on g(.) and > a universal constant independent of g()) and n,
55

---

G.4.1 Preliminary and Additional Notations
We first introduce notations and preliminary background. We consider the following stochastic gradient algorithms (SGAs)
2n+1 2n Ne En),
where 2n RR+ is the parameter vector, ?e is the effective learning rate, h(:..) RR RR+D Rd++ depend on the current parameter vector dn and the noise vector &n sampled from some distribution E(x,).
We also consider the Stochastic Differential Equation (SDE) of the following form:
dx, ;, ()dt 1)d,
where b IR++D RR+ - Rd+D the drift vector function and a RR+D R+ R(d+D)x(4+D) is the diffusion matrix function.
According the moment calculations in Corollary G.2 we set l/e - and
Lemma G.11
Lemma G.12 and Lemma G.13
b(x,,1)= = 0 (E()
(1 .c(V(E(()) -
a (X,,,)) ( (10 S(v)),
0)
0.
Next, we are going to define the one giant step change the parameter, both for SGAs and SDE.
X (1Rgr)) = ( g (0) ullaa ') T
E Rd+D (m)) = X ((n+1) (n (ngg)))
(m)) X(n+)mn X (n
((m) = b ((x (mRzo)),)
((x) = o (x (n
We now give a lemma to give the approximation of the first, second, and higher-order moment change of the SDE.
Lemma G.15. There exists a positive constant co independent of Th and such that for all Ç E T, holds for all 1 & i S d that
EJA.(C, - mc)
S come.
E (JA.(S.n) - ne Σ o.i(C)ou,(())
$ coTee
II A. (c,
< com
Proof. (i) By Lemma G.14. the first half solution ((t) X X(t) of Equation (8) stays in the manifold
almost surely when ((0) E r. (ii) We assume that C E C so b, o C (ii) We know that I' is
compact by Assumption 3.4 Then we can directly apply Lemma B.3 in Malladi et al. 2002 and
Lemma 26 in LinLi et al. 2019].
Lemma G.16 (Adaption of Lemma 1.41 in Gu et al. 2023bl). Given drift term and diffusion term b.o G and Lipschitz. Let E [0, 2] and g E G Then
[s, T., we can define:
u(æ, 8, = Ex, (X.))
where Px(x 8,1) denotes the distribution of X. with the initial condition X(s) ac. Then s,t, E 6 uniformly in s.t.
56

---

G.4.2 Proof of the Approximation for Slow SDE of AGMs
For the giant step constant 3 E (0, 0.5), we define several quantities ai 1.5228 (1, 11,),
1-8
02 1 1-3 (1.2), 03
2-28 In this part, we will show that only 1.5-1.59 1.5, and a4 1-3 1-8
dj and 02 would impact the error bound in our approximation theorem.
The following lemma captures the difference between the SDEs and the AGMs' first and second moment changes, as a key step control the approximation error, utilizing the moment calculation results from the last section.
Lemma G.17. If (10 6(0)
Ol )),, then it holds for all S n I [T] [T/mel and
1 S i that
E
(n)
(1)
Eg
B
Ne (log
1 )6 + ne (log
)0)
ne
T7c
B[A
(n)
A (m)
3
A
(m)
E0
e(mBor))
le (log
1 ee (log 1)
17c
6
E
(m)
g(nRgg))
I cG
2ai (log
y26
17c
6
E
(m)
g(nRero)
S Gi 201
(log
_20
1/c
where 0j and b constants independent of me and g.
Proof. According to Appendix G.2
we have that
11
A (n)
Eg(n
0(13333))
We can further use Corollary G.2
Lemma G.ll
Lemma G.12 and Lemma G.13
which gives
E[A
E
(m)
nob, (x)
1e (log 1 _ )o + ne (log 1 )))) S 02
(9)
ne
E
A
(n)
A
(n)
nc
(n) (n) i Li.
S 02
No (log ne
yb + nea (log 1 )0) 10)
(10)
17c
1=1
(n)
(log 26
(11)
ne
Notice that the above equations uses < ag and 02 < for all B E 0.5). These three equations and Lemma G.13lgiv the Lemma.
Lemma G.18. For test function E C and we define (x,) = u(xr, Lnfe mm)e)
If [[0(0) (x)))
(000 ||2 Olv /n log H)) then for all S 1, and
1 < n S [T],,, holds that
1E[441, (X ((Rx)) + A)) 4141, (X (1 + A0) X (1Ren))) s Cg.2(1 + n2²) log(
1 )
T7c
where Cg.3 some positive constant independent of ne but can depend on g.
Proof. Given E C by Lemma G.16 we have ur,n (x) E C for all and n. Which is to say that there exists function Q(-) E G. such that the partial derivative of un (X) with respect to I, up
57

---

to the third order is bounded by Q(x:). By the law of total expectation and triangle inequality,
E 1 (X (1ux)))
+ A (0) u1+1m (X (1 +
E 2[u++1, (LRgrp) + (1) W1+1m & (1Rgp))
(()p 000
c((
1 100 B u441m (X (1 + A
x (1 E (
12
n 100 E u+11 { (1 + + (1(Baa c
For I2 and I3, due to the compactness of I and v = R1 from Assumption 3.3 ((x) can be bounded for some constant 9911 independent of ne but could depend on test function g. Hence, we have that I2 + 13 S 100
Using the triangle inequality, we first decompose 11 into several terms as
@us.n 0x,
(Bar)) (4 - =.") X (1 (RRs)), 20
e(UR
1 <
11.1
1
Pugn (B(or))) (A!) - A!D
& (1Rarr), eo ell
E
ax,0x,
ISijEd
1,2
+ R + RR)
where the third order remainders R and R are
R
Σ
E
(Bee) AA 44)"
( ,0 eft
0x,0x,0Xx
1si.j.ked
R.
Σ
E
Unn Xx,,
(R)) + +aā0) )) (A (1Rg))) el c(l
where a E 1). Again, notice that the P is compact and ou 3 R1 thus we can bound the derivatives of (x) for any X as
0u141.m (X) S Cg.
i Putin Ox,0x,
Puttin (x) s CC
(x) < Cg.
(12)
0x,
0x,0x
For the term I1, and 1,,,, by applying Lemma G.17 we have that
y 1,2 s 2 C Cna + n (log -
de
I. S de, Co + n (log Te
Te
Next, we bound the remainders R and R. By Cauchy-Schwarz inequaaaatazzznuuulittt....
RI <
Σ
IE Junn (X(1) + a."))
X (L E ((l x
,X,0
(A (1) A (1) (00 (1Ron)) &n ((Bor)))
E
s '6 Co.crm log( )0,
13
n7c
58

---

where the last inequality uses Lemma G.17 and Equation (12)
Similarly, prove that there exists positive constant Cg such that
Cg,scime log( -
R
Combining the bounds for /, 12, and la gives the lemma.
G.5 Proof Theorem G.2
Finally, we ready to prove Theorem G.2
Proof of Theurem G.2
For 0 S n L,p J. we denote the random variable by L1,a such that
follows distribution Px (X (Uarr)) Me, m.ee). When we set / n, P(on, n X (()))))) and setting
= gives žo,n (nnee) Recall the previous definition that u(z,
Px (g (X),)
and define that Ti+1, == u1+1, (X (UR + A (L+ 1) nmee) 11+1, (X (1( + (1) (1
1) me)). Using the definition of Ti,n. we can rewrite the distance between AGMs and SDE measured by test function as
g g(x ((( - g( ( (mme)))
E g(an..)) g( e(
+ O(1y 100).
The above equation the law of total expectation and the definition of ô-good event &g (n in Definition G.I Then the Triangle inequality gives
E (( - ( co
c(nBpp)
E (g( - g(tr,n) e
+ O1))00)
Ti+1.
E e(mRRarod))
+ (1)
T4i. (1 e c(nR
+ O11)100)
E [ Te+I, (1 g ((n
c) c(n
+ 100)
1
S SnCg.aane + n (og))
1/c
1
TCg + n 1
log(
-
)6.
77c
where the second last inequality uses Lemma G.18 Recall that
1.5-28
1
B E (0,0055
1-
1-
Let 3 = 0.25, and we complete the proof.
H Proof of Theorems in Appendix B
H.1 Proof of Adam and AdamE's Implicit Biases with Label Noise
In this part, we give the proof of Theorem B.
Lemma B.] and Lemma B.2
Proof of Theorem B.1 Recall the SDE formula Equation and Lemma G.10
{ dc(t) = 04 s()) (C)())((()))))))) (() (G
PId,
du(t) e (V(E()) v) dt.
Plugging in the following:
59

---

The definition that P = -S0?2
Lemma G.3 For any < E I' and p.d matrix S, a (C) =,
The label noise condition: Σ(() = aV (()) for any < E T and some constant >> 0.
yields the final result:
{
dç(1) ; 5, (VC. dv(t) (V(>(()) v)
(13)
The above equation completes the proof.
For Lemma B.I and Lemma B.2 first present the following formal statements and give the corresponding proofs for each of them.
Lemma H.1 (Adam's Implicit Bias under Label Noise and e - 0). With the label noise condition, every fixed point of Equation for Adam with satisfies Vrtr (Diag( H H) where Vr stands for the gradient of Junction projected to the tangent space of T.
Proof of Lemma H.1
Consider the a fixed point (5","") of the ODE 13) must satisfy
S(v" ((")))0((((((((()""")")""
-,
(14)
and
V* = V(E((*))
(15)
First, we simplify the notation by denoting the following:
S == S(v)),
Pi 0 s( S H' V2C(())
Then Equation (14) becomes
S' C'"'''(
Since x*' V(E((") and E((*) aH' we have that 2" diag(E((")) adiag(H") Then S Diag( (adiag H )31/22 ) and we can rewrite ( as
d
3 [' 11 V(H)))
2 Va
V ((H
2 Vtr (Diag (Diag(H)))12 1/2) Vo
Therefore, Equation (14) is equivalent
Sth S'PiVtr (Diag( (Diag(Hr*)377 H222 0.
W.I.O.G., decompose P and H into block matrices
0. 0, P
2 .
0. 0. Him)
Pa Pi
H
where Pi-m H R (d-m) are full-rank matrices. Under this decomposition, the first m diagonal elements in (Diag( H/2 is O, and the first diagonal elements in S* is 0. Specifically,
0 Diag H -() Diag H2)m))
S
S
Then the constraint in Equation (14) can be reduced into
Vrtr (Diag(H+))) (/2) 0.
(16)
proving the theorem.
60

---

In the practical use, we usually set to an extremely small constant; for example. Py PyTorch's 's official documentation sets the default value of to 10 -8 Such small constant will hardly alter the form of the implicit bias, but to make our analysis more rigorous, we also derived the case for Adam with
< <> 0.
Lemma H.2 (Adam's Implicit Bias under Label Noise and e 0). With the label noise condition, every fixed point Eqqation ] Adam with e satisfies that
Vrtr Diag H) In (VaDiag( H 1/2 +c)) = 0. H)
va
Proof. Let (S", v*) be a fixed point of the ODE (33 (equivalently. of Equation (3) Then, as in the proof of Lemma H.1 it satisfies the stationarity constraints
S(v (" 0,
(17) (18)
2" = V==(())(())
Introduce the shorthand
S = S(v*),
Pi = (..) =
= (ZC(().
Under the label noise condition we have Σ((*) = aH hence by 18
= V(E(()) diag( E((")) adiag H
For Adam with > 0, the diagonal preconditioner is
s = -Diag
Diag Vddaag( (H) +
Vo" +
We now compute V [(()) Since S is diagonal, we only need to sum up the diagonal terms:
V'[((")[S- Σ 1
Vallis +e v(Hi,).
(19)
Define the scalar function
(x)) = Vax eln (Vax + E),
x 2
direct differentiation gives
v'(x) 2 2(Vax +6)
Therefore,
2 2v4.
Vax 199 and using the chain rule yields
Plugging this identity into
V'C((")[S)] 2 d Σ (HH) 2 v (HH
2 Vtr (i((U(Dig((H"))) (20)
Noting that & acts elementwise on the diagonal, we can also write
tr dh (Diag( H-)))
((DDin((H-)12 eln (VaDiag(HH))l
Substitute 20 into
2 Vtr (/DDiag(H')u eln (VaDiag H 1/2 + e)) 0.
P a
Since >> 0, the diagonal matrix S* has strictly positive diagonal entries and is therefore invertible, and we can cancel out. With arguments similar to those in Lemma H.I only the bottom-right part of H" (that corresponds the dimensions within I') are nonzero, and we obtain the following result:
Vrtr (VaDiag( /2 - In aannn(aaaan (VaDiag(H /2) + +e)) = O.
61

---

direct calculation gives
V Va ( Diag((H))) +1)) O, a
Vrtr (Digg(H)) ( In
which proves the claim.
Lemma H.3 (AdamE's Implicit Bias under Label Noise and 0). With the label noise
dition, every fixed point of Equation (3) for AdamE-X with À E [0, and e = O satisfies
Vpir (Diag(H)²* =0.
Proof of Lemma H..3
The proof is similar to the proof of Lemma H.l but when calculating S we (--))x) instead of S Diag( (odlag H "))²/2 ). which gives
have 5 -Diag( (odiagg H
Σ
(aHij)A V(H))) = (1 N)a Σ j=l
(((H)))*>>
1=1
Vtr (Diag(
(1
d)ed
This leads to the constraint - Vrt (Diag H*) = in the same way as Equation (16) which completes the proof.
H.2 Proof of Lemma B.3
Proof. We only prove the second argument in Lemma B.3 with any 6p E (0, 1], since taking ép = 0.5 yields the first argument. First, recall that the minimizer manifold T is defined
{ (z,, u 02
02 yo Vi E [m]]
So if any (*) belongs I, and another 0 (+) satisfies that u 0.2 u, U
02
for any
i E [d]. then 0 also belongs to r.
Next, we derive the explicit expression of the Hessian matrix when 0 l::
x
S. 2
(.0 T
+ ((zi, 02 022) yi) Diag(=)
O
)
-Diag(z)
71
2,00 (2.00)
2:0 (22,0w)
Hence, we have that
tr tr(Diag( H ()) x
(1us 2 + 200))
i=1
and ||002 02 02 ||
==1 |u
oo Let (0, 1]. and we recall that our goal to prove
that given the following condition
E arg min tr(Diag (H) = arg min d'el
(Iug 260 + |0i 200)
(21)
d'el
holds that arg ming ||
First, prove that u, 0 V v; 0 holds for any E [d]. Assume for the contrary that there exists
some such that u, * 0 and v, 전 then construct another reference point Ø (*) by letting ù and u agree on all indices other than i, and that
ú, v ug of, U. = 0,
if |uyl 2 |val,
(22)
/ 12 u
otherwise.
(23)
62

---

With this construction, u 802 =u
u0
v
1 E Т. One can observe that
ju, 20g
[6
20o
< [ 200 + |0
which contradicts the condition in Equation (21)
Also, we prove this by contradiction. Now assume 0 & arg min /202 02 There must exist some E I such that
Now we are ready to prove 0 E arg minorer
@ eoo
7 02
02
1 02
02
W.L.O.G., one can assume that for any minimizer that preserves | i
E
[d]
either ū; or 0 0. else we can construct another Equation (22) and Equation (23) However, given the
condition u, 0 V u, 0, we have that
Σ
u
2000
Uu 2e0 + |0//200 200
and Di=1 lii 00700
|
2e0 + |0,|200, which indicates that Σi= [ū 22e + 10
220
11=1 [u:] 200 + |0:|200,
220, a contradiction.
I Regularizers under label noise for AGMs in Table 1
In this section, we provide additional discussions on the regularizer for the AGM optimizers in Table besides Adam, AdamE (refer to Apppeddxx H). and Shampoo (refer to AolAppendix J).
SGD. Under label noise, the implicit bias of SGD has been extensively studied by previous works; As discussed in [Section 2, approaches such fixed point analysis [Blanc al. 2020], slow SDE (Li et al. 20216] and implicit gradient regularization [Barrett and Dherin 2020] all agree on the result that SGD implicitly regularizes tr( H) on the minimizer manifold. Our work provides new insight on the implicit bias of SGD by comparing with that of Adam. Specifically, SGD treats each direction equally which results in rotation invariant (r( H)) the implicit regularizer, while Adam has the second-order momentum as a denominator, so Adam regularizes the entries with small gradients relatively faster, as is indicated in its implicit bias tr tr (Diagg
RMSProp. The RMSProp optimizer Hinton et al. [012 can be seen special case of Adam, where 31 One can observe that 31 does not appear in the slow SDE system, which implies that as long as 1 - 3i is of constant order, the choice of 81 has nothing to do with the dynamics of Adam on the minimizer manifold. The intuition is that, after the iteration approaches the manifold, the gradient (L(s)) moves very slowly as A proceeds. Since the momentum only captures O (log 1/1) past steps, the different between momentum and the gradient at that step becomes negligible. Therefore, RMSProp possesses an implicit bias identical to Adam: tr tr(( (r(iiigH)) H) 1/2).
Adam-mini and Adalayer. Adam-mini [hhan et al.) 2025) and Adalayer (Zhao et al. 2025 belong the same kind of variant of Adam that partitions the parameters. In Adam-mini the partitions blocks, and in Adalayer the partitions are layers. In the sequel, we provide a brief derivation of the implicit bias of "Partitioned Adam", which is applicable to any kind of optimizer whose functions V and S can be expressed in the form of
V(M), = 1
Σ
[Ba(O]
jEBrco
S(v) = Diag(1/(v)))) + e))
where B = (B,, .z.... BN] is partition of [d], and for each E [d], T(1) denotes the index of the set containing i, i.e. E Ba(s). We derive the case for O.
Recall from the proof of Lemma H.T that the gradient of the implicit regularizer being minimized on manifold can be expressed as
82 (V)) = V ((S H V (S) H
(24)
63

---

In our case S is diagonal, so we can calculate the contribution of each set in the partition, and add them up. Next we focus on single set, and re-index it as 21. ... G} without loss of generality. In this set trH/ / is used as a shared second-order momentum, so we have
trH G
1/2
s =
la
Combining with P S1/2 gives us
(S, H) (PHP) VG trH.
For the second term, we again denote hj == H and we further denote t: :: tr/HH 1 EE=1 hj-
(S) = Σ 112) h3
= Σ v (t) h1
_ --12
1
2
= GVV - - 1-12
1
-V (1=²)= = -
Plugging into (24) gives the implicit bias contributed by this set as /G trH. Finally, summing up all the sets, we conclude the overall implicit bias as
Σ
RE(N]
Here, H8, means the submatrix of H if we restrict the rows and columns to B,
J Shampoo Optimizer as an AGM
J.1 A brief introduction to Shampoo
Shampoo, unlike most conventional stochastic first-order optimization methods, utilizes the fact that many model parameters in practice tensor-like, and lead to faster convergence in optimization. In this paper, we consider the case where the parameter is a matrix (2-dimensional tensor) with shape dj d2. In this case, the Shampoo algorithm be expressed follows:
Algorithm 1 Shampoo with matrix like parameters
Require: horizon K, learning rate > 0, stabilizing constant e >> 0.
Initialize @ Xda;;;; S1,0 1ddixd, $2,0 xd
2: for k = 1 to K do
3: Receive loss function la IRu xda ->
4: G t (Qx) E Rdixda
5: Update S1,++1 with Si, and G
6: Update S2 with S2,1 and G. Gk
7: @x+1 < O n (S1 cIa,) G (S2 + cld.)
8:
end for
Here, is constant that equals 1 /4 when Shampoo was originally proposed by Gupta et al. 2018 8
while later works Anil et al. 2020, Shi et al., 2023 Morwani al. 2024 suggests switching to
1/2. We incorporate the constant here, being consistent with Gupta et al. 2018 and the practical need to stabilize training. There are multiple choices for the update rule in Lines 5 and 6. Originally, Gupta et al. [2018) simply sum up the past terms to maintain S and Sz. Later in practice [Morwani et al. 2024 Lin et al. 2025), this was replaced by an Exponential Moving Average (EMA):
S S = 32 S a + G.G.T.
S = B3 S2, + (1 - (2) G. G
(25)
64

---

We adopt the update rule 25 for Lines and 6 in Algorithm 1, which aligns with practical imple- mentations and naturally fits within our AGM framework: We use 1/2 to stick with the latest result, but the following analysis actually holds for any constant À 2 0.
J.2 Shampoo under the AGM Framework
In Shampoo the parameter takes matrix form, while our AGM framework requires a vector, need to define some reshaping rules to view Shampoo seamlessly as an AGM.
First, we define the vectorization of matrices vec(-) R2*0 + RM as
X I, where Fig+j = - A 0 <
We introduce the Kronecker product, a matrix operation that generalizes the outer product matrices. Given any A R2x4 and B E IR* their Kronecker product A @ B E Ruexx* is defined as the block matrix obtained by multiplying each entry of by the entire matrix B. It satisfies several useful properties:
vea (AX (B A) veo for any X R
(26)
A'' & B¹) (A B); for any A E S B DES4,, R.
(27)
where 4 S4+ denote the set of positive definite matrices in TRX and R respectively. For any step & E [0, K - 1]. we define 0₂ := vec ()) and gl := vec GG)). Let = d1d2, then OR.gr E IR4. The Shampoo update (Line Algorithm 1) can now be rewritten
0k+1 0x - ((5 + cla,) @ (S1 + +Ia.))
gN.
(28)
Next, for any M Raxd we introduce a new indexing scheme that represents M as a 4-dimensional tensor by decomposing each row and column index into two sub-indices. Specifically, for all S i.j d1 and 0 < 1, d2, we use
M.,.r.j
as an alternative representation of
Midg+1, jdgtr
Then we define two functions Vz Raxd pd xd, VR Rdddd + d2 as:
[V) = Σ Mi,e,
(VR(M))s = Σ Ma3d..
Note that for any step As E [1,
(V Dl.s Σ [gn Iississ
LIGH [G
(29)
[VR(9ngT 1)., = - [g J..ij. L. = [ G
(30)
Now are ready to rewrite the Shampoo optimizer in the AGM form. Definition J.1 (Shampoo, written in the AGM form).
V == Vi + (1 (2
(31)
Ox+1 = 02 (VS((++1))r.
(32)
Here, each VA E Rdixds x iddaxda is tuple of two matrices with shape d1 x d1 and de x d2 respectively: The functions and S are defined as V ddd Red xd, x ddd
M + (x(M))
+ (12 + cla)) TT & (Vi + =aa.)) -1/2
S Rddxdd xd;
R xd2 ddxd
With
29 and 30
it straightforward to verify that recovers 25 and that V = (Si,k, S2,1);
Hence 32
recovers
(28) as well.
65

---

J.3 Vectorized AGM form
is worth noting that Definition J.I we slightly generalized the definition of and V so that Ve not a plain vector now, but tuple of two matrices. This is only for a simplification of the expression of functions S and V which makes them easy to understand. of Shampoo that is slightly more complicated, but rigorously fit in the shape required by AGM, and we will use this form the next subsection to analyze why Shampoo's 's bias on the manifold cannot be written using explicit regularizer.
this subsection, we establish form
For specific shapes of matrices: d1 d1 and d2 d2, define functions that inverse the vectorization effect:
mal RR Rdixd,
xd.
(w)ls Vidi 6A i.j < di,
matR R RR Rdaxds,
( Vida+j 44 S i.j d2-
For any vector v E R we write (Uz, UR) and v interchangeably so that vr, represent the first di entries and vn represent the rest. Now we present a vectorized version of Definition J.l Definition J.2 (Shampoo, written in the vectorized AGM form).
Uk+1 == 8u + (1 V
0 0x+1 == 02 ( 1)gk-
Here, each 2s R.D where D dt dz, and the functions and S are defined as
V Rddd
R.D.
M + (vec (M)) vec (VR(M)))
(matz(vl) eld.)) -1/2
S RR dxd
u + (matt eld,)
J.4 Discussion on Shampoo's Implicit Bias under Label Noise
Recall that for all AGMs, under label noise, Equutiin (14) and Equation (15) hold, and we adopt the notations S* Pl H" in Lemma H.T to write them
((")S" =0 0" =V=V(=(("))
Here we can expand S as
S S(v*)
((matR(Vi)) + +el,)) (matz(vi) + eld.)) -1/2
((Va(E((")) + +laa)) @ (Vz(E((")) + +ela.)) -1/2
and we further denote
= = Vz(E((*)) + cla = Vn(E((")) + ela
We define function
A: Rd Rd as
A(C) +*
Note that
A - +'
(33)
We provide theoretical evidence that Shampoo may not admit any explicit regularizer under label noise by showing that the function A is not conservative vector field, thus no potential function v exists such that V A.
To see it clearer, if such a function v exists, then with the same techniques as Lemma H.1 we can simplify [14) into
Vrob((*) 0,
66

---

so the regularizer is exactly (5(5),, and vice versa.
Unfortunately, in simple case, where H is assumed to be diagonal (for example, the scenario of diagonal net), the potential function y does not exist. In this case, we can assume that
H Diag(An, , da);
where X 2 0 for any S S d. Consequently,
E(()) aH" = aaDig((A, A2.....,, , Aa).
Now we have
Σ = V2((((())) +eld
do
dg
de
= aDiag
A1tidi
2
Ad Add
+ +
1=0
1=0
1=0
d2-1
P2.. rd,) E Rdxxd
where rj =
Ajtids 6
Σn Vr(E((*)) eIdd
d1
d1-1
d1-1
aDiag
A1+ada
A2tid2
Adztidz
+ cl
d1 -1
= Diag((1, ,,,,,,,,, Rddxds
where 1j = a
Ajtida +
1=0
Therefore, the preconditioner matrix S in Equation (33) can written as
S = [E @ Σi)11/2 1/2
Diag 1/7 (ER) 1/2 1//,,11/1//
1/2 (ER) 1/2 ,...,,",, 1 ((E))-11
One can straightforwarddy verify that the curl of at 5":
x A(5") V (((")
is nonzero. By the Stokes- Stokes-Cartan Cartan theorem (Theorem 16.1 in Lee 2012 potential function vb such that V = A. Therefore, we argue that in general, the regularization effect of Shampoo under case is counterexample.
there does not exist
label noise cannot be reduced an explicit regularizer, for which the diagonal
67


Reviews:
[{"review_id": "eeRXPzsQQC", "reviewer": "Reviewer_8TUV", "summary": "The paper theoretically analyzes the implicit bias of Adam. One of the main results shows that in the setting with label noise, Adam and SGD minimize different things: Adam minimizes $\\mathrm{tr}(\\mathrm{Diag}(H)^{1/2})$ and (as has been previously known) SGD minimizes $\\mathrm{tr}(H)$. Furthermore, the authors propose a tunable version of Adam which minimizes $\\mathrm{tr}(\\mathrm{Diag}(H)^{\\lambda})$ for any choice of $\\lambda \\in (0,1)$. This allows more control over the desired implicit bias of the algorithm.", "questions": "* In the experiments, all version of AdamE either behave more or less like SGD or like Adam. Can you narrow down at which $\\lambda$ the change happens? Or, if we tried many more intermediate values of $\\lambda$, we would see a gradual shift requiring more and more training data as $\\lambda$ decreases? If it is sudden, do you have an explanation for this phase transition?\n* I assume \"Aappendix\" in line 302 should say Appendix A?", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Unfortunately, I could not check all the math, but the paper appears very interesting to me.\n\nStrengths:\n* Adam is considered far less in theoretical studies than GD is despite the fact that Adam is, used more in many practical settings. This paper adds to the understanding of Adam.\n* The tools presented in the paper capture other adaptive gradient methods and should provide useful for further works in the future. It is almost more like a framework than a single result.\n* The authors make an effort to try to illustrate what the difference in behaviour between SGD and Adam mean intuitively.\n* Proving Adam's bias towards sparsity in the presence of label noise is very interesting.\n\nWeakness:\n* There are a few assumptions (but I found them mostly mild).\n* The analysis only works for specific parameter settings: $\\beta_2$ and the learning rate have to follow a specific relation.", "was_revised": true, "final_justification": "There seems to be broad agreement among reviewers on the value of the work, the authors engaged with my questions and I am happy to keep my positive score.", "rebuttal": "### Response to Reviewer 8TUV\n\nDear reviewer `8TUV`,\n\nWe sincerely thank you for your interest in the paper and your appreciation of the slow SDE of Adam. We are pleased to hear that the paper's results, particularly the analysis of Adam's bias towards sparsity in the presence of label noise, are considered “very interesting.” Moreover, we are grateful for the acknowledgment that the tools and framework presented could be valuable for future work in adaptive gradient methods. Below, we address the specific concerns and questions raised in the review.\n\n**Q1: In the experiments, all versions of AdamE either behave more or less like SGD or like Adam. Can you narrow down at which $\\lambda$ the change happens? Or, if we tried many more intermediate values of $\\lambda$, we would see a gradual shift requiring more and more training data as $\\lambda$ decreases? If it is sudden, do you have an explanation for this phase transition?**\n\n**A:** In our experiments, the **transition is sudden, and happens very close to zero, between $\\lambda \\in (0.01, 0.001)$**. To provide an insight into this, note that **the change of minimizer may not be smooth when we gradually change the norm being regularized**. \n\nHere we provide a toy case for understanding the sudden phase transition caused by the regularization parameter  $\\lambda$. Recall that AdamE-$\\lambda$  minimizes $L_p$-norm where $p = 1 - \\lambda$, so the phase transition happens around $p = 1$. We consider a setting where we have two parameters  $x, y$  and the manifold being  $x+y=1$. A point with minimal  $L_1$  norm on this manifold could be any point on the segment between  $(0, 1)$ and  $(1, 0)$, but for any  $p < 1$ , the  $L_p$  minimizers on the manifold suddenly shrink to only two points, $(0, 1)$ and $(1, 0)$.\n\n**Q2: I assume \"Aappendix\" in line 302 should say Appendix A?**\n\n**A:** Thank you for pointing it out! After some reordering, this part of the proof actually exists in Appendix G in the supplementary material, and we will fix this in the camera-ready version.\n\n**W1: There are a few assumptions (but I found them mostly mild).**\n\n**A:** We do make a few assumptions, such as the smoothness of the loss function and the existence of a minimizer manifold, which we agree are mild, as acknowledged by the reviewer. We also look forward to further refining and weakening these assumptions in future work. \n\n**W2: The analysis only works for specific parameter settings: $\\beta_2$ and the learning rate have to follow a specific relation.**\n\n**A:** Thank you for your careful reading and thoughtful feedback. \n\n- In this work, as the first step towards uncovering the implicit bias of Adam under a general setting, we analyze Adam under the case where $1-\\beta_2 = O(\\eta^2)$, namely the  $2$-scheme. \n- For any $p$-scheme with $p>2$, the second-order moment will not vary significantly throughout the whole implicit bias phase, so Adam will simply become a pre-conditioned version of SGD, and the conclusion can be simply transferred from that of SGD. \n- Although we did not cover the case where $p<2$, the SDE derived by us is already useful, since it allows for the first time the analysis of Adam's implicit regularizer (Section 5.1) and the anticipation of Adam's advantage/disadvantage over SGD on sparse linear regression/matrix factorization. Subsequent works can build upon ours and derive a tighter bound, which is beyond the scope of this paper."}, {"review_id": "HCOxNu2HF7", "reviewer": "Reviewer_GRmZ", "summary": "This paper investigates the behavior of the Adam optimizer under the assumption that the set of minimizers forms a manifold. By analyzing the corresponding slow stochastic differential equation (SDE), the authors demonstrate that Adam  promotes sparsity in the solution.", "questions": "See the weaknesses above.\n\nMoreover, Diag(H)^{1/2} is Diag(H^{1/2}) or (Diag(H))^{1/2}?", "limitations": "N/A", "rating": 4, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "I'm not familiar with the field. Here are some strengths and weaknesses:\n\nStrengths:\n\nThis paper generalizes the slow SDE to more general adaptive optimizers, like Adam. The result also tells us the difference between SGD and Adam minimizers they find.\n\nWeaknesses:\n\n1.The assumption implies that the trajectory x(t) of the optimizer converges to a single point, which can be quite strong. In practice, especially in non-convex settings or under noisy dynamics (like SDEs), only subsequence convergence may be guaranteed.\n\n2. Why Assume the Minimizer Lies on a Manifold? This is often based on empirical observation as said in the paper. However, from a theoretical standpoint, assuming the set of minimizers is a manifold requires strong structural properties.\n\n\nminor comments:\nLine 159: better to use the terminology: in the normal space.\n\nline 202: preconditioner flow: lacks blanket.\n\nDefinition 4.3: Typo Correct: given → Given.", "was_revised": true, "final_justification": "The issues are solved. I have raised the score.", "rebuttal": "### Response to Reviewer GRmZ\n\nDear reviewer `GRmZ`,\n\nWe appreciate your thoughtful and constructive feedback. We have carefully considered all comments and provide our point-by-point responses below. We believe these responses will address your concerns and improve your understanding of the manuscript.\n\n**Q1: The assumption implies that the trajectory x(t) of the optimizer converges to a single point, which can be quite strong. In practice, especially in non-convex settings or under noisy dynamics (like SDEs), only subsequence convergence may be guaranteed.**\n\n**A:** **The iteration does not necessarily converge to some fixed point under our assumptions.**\n* Our assumptions regarding the loss function include: the smoothness (Assumption 3.1) and compactness (Assumption 4.2) of the manifold that the optimizer converges to; and the smoothness of the loss function (Assumption C.1, C.2). Note that we don't require any assumption such as convexity that will result in a single global optimum.\n* We could not find the notation “x(t)” used in the review besides that in our definition of gradient flow projection (Line 154, Line 202). Here, x(t) is defined as an ODE, which does not have \"noisy dynamics\". We define the gradient flow projection only when the limit point exists, so this notion is well-defined.\n* If “x(t)” actually means the SDE trajectory “$X(t)$”, then it is true that $X(t)$ may converge to a single point in some cases. However, even this is NOT implied by our assumptions. Some other assumptions are needed to ensure this convergence. Even if we add more assumptions to make $X(t)$ converge to a single point, Adam's trajectory does NOT converge to a single point. This is because $X(t)$ only describes the continuous approximation of Adam projected onto the minimizer manifold. The original trajectory of Adam can be seen as $X(t)$ plus some perturbation of scale $O(\\eta^{1/4})$ (Theorem 4.1), which may not decay to zero when $t \\to \\infty$.\n\n**Q2: Why Assume the Minimizer Lies on a Manifold? This is often based on empirical observation, as said in the paper. However, from a theoretical standpoint, assuming the set of minimizers is a manifold requires strong structural properties.**\n\n**A:** The manifold assumption is still reasonable from the theoretical perspective for the following reasons: \n- **Under an over-parametrized setting, the minimizer manifold provably exists** with mild assumptions. Specifically, with parameter dimension $d$ and dataset size $n \\ll d$, define the set of global minimizers as $\\Gamma$. Two commonly used assumptions are: 1. all global minimizers interpolate the dataset (i.e. $f_i(\\theta) = y_i, \\forall i \\in [n]$ for all $\\theta \\in \\Gamma$), and 2. the Jabobian matrix $[\\nabla f_1 (\\theta),\\nabla f_2 (\\theta),\\cdots,\\nabla f_n(\\theta)]$ is of full rank $n$ for all $\\theta \\in \\Gamma$. With these two assumptions, **the preimage theorem** implies that $\\Gamma$ is a $(d-n)$-dimensional manifold. A similar discussion can be found in Section 3.1 of [4].\n- The tendency for minimizers to form connected sets, namely the **\"mode connectivity\"** phenomenon, has been widely observed in practice [1,3] and has been demonstrated theoretically in some notable settings [2]. \n- **The manifold assumption is a standard practice** used by many theoretical papers in the literature [4,5,6,7], which is a close approximation of mode connectivity for the tractability in theoretical analysis.\n\n**Q3: Does $Diag(H)^{1/2}$ mean $Diag(H^{1/2})$ or $(Diag(H))^{1/2}$?** \n\nA: $Diag(H)^{1/2}$ means $(Diag(H))^{1/2}$. That is, first we keep only the diagonal entries in $H$, and then take the square root. \n\n---\n#### Response to Minor Comments\n\nAlso, we thank you for your careful reading and helpful suggestions. We will incorporate all of these corrections into the revised manuscript. Specifically:\n- Line 159: We will adopt the suggested terminology \"in the normal space.\"\n- Line 202: We will add the missing space to correct \"preconditionerflow\" to \"preconditioner flow.\"\n- Definition 4.3: We will correct the typo (\"given\" → \"Given\").\n \n[1] Loss surfaces, mode connectivity, and fast ensembling of dnns.\n\n[2] Understanding warmup-stable-decay learning rates: A river valley loss landscape perspective.\n\n[3] Explaining landscape connectivity of low-cost solutions for multilayer nets.\n\n[4] What Happens after SGD Reaches Zero Loss? --A Mathematical Framework.\n\n[5] Why (and when) does local sgd generalize better than sgd?\n\n[6] Convergence rates for the stochastic gradient descent method for non-convex objective functions.\n\n[7] Understanding gradient descent on the edge of stability in deep learning."}, {"review_id": "1pedQbTHlI", "reviewer": "Reviewer_cpiw", "summary": "This paper studies the Adam optimizer and shows that it reduces a specific form of sharpness that leads to qualitatively different solutions from SGD. The paper shows that when using Adam to train over-parameterized models with label noise, Adam minimizes $tr(Diag(H)^{1/2})$ a measure of \"sharpness\" which differs from the $tr(H)$ that is minimized with SGD. Finally, they show that for the specific case of solving sparse linear regression with diagonal linear networks Adam achieves better sparsity and generalization than SGD.", "questions": "None.", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths:\n- The paper is written clearly and the results bring new insights into how the implicit bias of Adam differs from SGD.\n- The paper provides concrete examples that illustrate the qualitative difference between the kinds of solutions Adam learns vs. SGD (such as Fig. 1 and the sparse linear regression example).\n- The paper introduces the AdamE-$\\lambda$ optimizer to directly tune the implicit bias of Adam.\n- The theoretical contributions apply broadly to many AGM and is not specific to Adam.\n\n## Weaknesses:\n- The entire Appendix is missing which makes it impossible to verify the correctness of the proofs.\n- Adam is typically used to train ReLU neural networks which are non-smooth and non-convex. It is not clear whether the analysis here translates to this more realistic setting.", "was_revised": true, "final_justification": "My main concern was that there was no appendix and it was unclear whether their results hold for ReLU based networks. Both of these have been resolved in the rebuttal process.", "rebuttal": "### Response to Reviewer cpiw\n\nDear reviewer `cpiw`,\n\nWe thank you for your feedback and for acknowledging that our paper “is clearly written and brings new insights” into Adam's implicit bias. We believe the reviewer's main concerns stem from not being able to correctly locate the appendix of our paper. We address the two key weaknesses below and are confident that these clarifications, along with the supplementary material, resolve the primary reasons for the borderline score.\n\n**Q1: The entire Appendix is missing?**\n\n**A:** We followed the NeurIPS traditions to put our appendix to the supplementary material. This can be downloaded through the 'zip' icon inside the *Supplementary Material* section in the OpenReview submission page, and the full paper `adam_sde_nips.pdf` can be found inside. We included the full proof, a formal restatement of theorems, an illustration of the difference between Slow SDE and conventional SDE, an additional matrix factorization experiment, and the code for the diagonal net experiment in the supplementary material.\n\n**Q2: Adam is typically used to train ReLU neural networks which are non-smooth and non-convex, how do your analysis transfer to realistic setting.**\n\n**A:** Actually, **our analysis is aligned with many realistic settings**:\n- First, **frontier LLMs have largely switched from ReLU activations to smooth activation functions**, such as SwiGLU, Swish, softplus, and GeLU, to improve training stability. Specifically, many mainstream LLMs such as Llama, Qwen, and DeepSeek series all use SwiGLU as the activation function. Therefore, it is reasonable to assume smoothness in our analysis. \n- Also, we would like to clarify that **our paper does not require either global or local convexity of the loss function**. Our main results (see full statements of our main results in Appendix C) focus on the dynamics near the minimizer manifold, where the local landscape may not be convex (but it can be proved to satisfy $\\mu$-PL property, a much weaker condition)."}, {"review_id": "jmPvCca4CG", "reviewer": "Reviewer_9RWz", "summary": "To analyze the optimization trajectories of Adam after reaching the minimizer manifold, the authors propose the *slow sde* framework, which applies to a general class of adaptive optimizers, including Adam. Authors prove that such a continuous SDE solution can approximate the iterates of Adam within an error lower than $O(\\eta^{1/4})$. In addition, they also provide a convergence guarantee for Adam to reach the minimizer manifold. Based on such an SDE approximation, the authors analyze the fixed points of the SDE and provide the implicit bias of Adam and AdamE-$\\lambda$, respectively.", "questions": "Could the authors explain more about ''slow sde\" for both the SGD and Adam? For me, the conventional SDE is clear. However, the \"slow sde\" is relatively hard to follow. It's difficult to understand why the trajectories of SGD can be characterized in this form. What do the diffusion term and the drift term indicate? Why is the current updating direction of SGD (from my understanding, this should depend on its local information) related to its convergence point (as reflected in the definition of $\\Phi$)? \n\nBesides, could you please explain how you deal with the second term in line 306 of the proof sketch? Since the definition of the calculation $\\nabla (\\cdot) [H]$ is defined for a Hessian matrix, it's a little hard for me to follow this calculation.\n\nIn addition, although you have mentioned that the first momentum would not affect the approximation result, I wonder whether you can also add it as a part of your \"slow sde\" system? This is not a major concern, just for better illustration.", "limitations": "No negative societal impact. Potential technical weaknesses have been discussed.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n1. The analysis regarding the continuous approximation of optimizers plays a key role in helping to understand their behavior. Therefore, such a result is of great significance, especially given that there are few related works for Adam due to its complicated updating rules compared to SGD.\n2. The theoretical findings and illustration of this paper are relatively complete and solid. Almost all the theoretical conclusions are established with convincing evidence and clear illustrations. For example, the authors provide both the convergence guarantee to reach the minimizer manifold, then the SDE approximation inner the neighborhood of such a minimizer manifold. I would, however, suggest that authors consider relocating subsection 4.3.1, the convergence of Adam, to the beginning of Section 4. \n\n**Weakness:**\nSome of the following points might be more suitable to state in the question part, but I still state them here since I'm not very familiar with the SDE and differential geometry. Due to the same reason, I have set my confidence level of this rating to 2. Authors could correct me if I have any misunderstandings regarding their results.\n\n1. The requirement that $1-\\beta_2$ is comparable with $\\eta^2$, coupled with the $O(\\eta^{1/4})$  approximation error established in Theorem 4.1, implies that the \"slow ODE\" presented in this paper might not adequately capture Adam's trajectories. For example, taking $\\beta_2=0.9999$, a value already larger than common settings, still yields a convergence guarantee on the scale of $1/3$.  Consequently, this approximation might not provide a strong guarantee. Is such a requirement raised as you use the $\\eta^2$ as the interpolation scale to establish the ode form of $v_t$? Similarly, I wonder whether there are some typos in the formula of Theorem 4.1 of the term $\\eta^2$. I also wonder why there exists a factor $\\sqrt\\eta$ before the diffusion term of the conventional SDE in Line 129. From my perspective, the noise of the stochastic gradient might be uncorrelated with the learning rate. \n\n2.  While the conclusion of the implicit bias regarding the trace of Hessian is convincing to me, some explanations of such an result would be much better. For example, what are the connections or comparisons between this result and the direct conclusion that Adam can maximize the $\\ell_\\infty$-margin in a recent work [1]. At least to me, the conclusion in [1] is easier to understand its implication. I know you have provided an example in subsection 5.2. However, its extremely specific settings ($z_i \\in \\\\{1, -1\\\\}$) make me feel that the conclusion related to the implicit regularization of different forms of norm lacks sufficient generalization. In addition, I suggest that the authors could design some experiments to directly support their main theorems on the trace of the Hessian matrix. \n\n[1] Zhang et al. The implicit bias of Adam on separable data. NeurIPS.", "was_revised": true, "final_justification": "Most of my issues have been resolved. I believe this presents good work!", "rebuttal": "# Response to Reviewer 9RWz \n\nDear Reviewer `9RWz`,\n\nWe sincerely appreciate your acknowledgment that our continuous approximation result “is of great significance,” and that our theoretical findings “are relatively complete and solid.” Your comments are invaluable, and we hope this response addresses your concerns.\n\n---\n\n**W1 part1: The requirement that $1-\\beta_2$ is comparable with $\\eta^2$ makes the $O(\\eta^{1/4})$ approximation error not very tight.** \n\n**A:** We thank the reviewer for raising this point about the approximation tightness.\n- Our main contribution is to correctly derive Adam's slow SDE for the first time, which already involves highly non-trivial technical challenges (Section 4.3). This SDE approximation is useful, as it enables the analysis of Adam's implicit regularizer (Section 5.1) and helps anticipate Adam's advantage or disadvantage relative to SGD on problems such as sparse linear regression and matrix factorization. While further tightening the bound would be valuable, we view this as an important direction for future work.\n- We also note that our current error bound is already the best known: even the state-of-the-art analysis of SGD [1] achieves the same order $O(\\eta^{1/4})$. Other works on SGD can only establish an asymptotic convergence of SGD to its slow SDE in distribution, without explicit error bounds [2,3].\n\n**W1 part2: Whether there are some typos in the formula of Theorem 4.1 of the term $\\eta^2$**\n\n**A.**  Thank you for your feedback. **We confirm that is not a typo actually.** This theorem shows that, in the small  $\\eta$  regime, once Adam approaches the minimizer manifold, its long-horizon behavior within $\\tilde{O}(\\frac{1}{\\eta^2})$ steps can be well approximated by the SDE defined in Equation (2). More specifically, in an SDE, a continuous time duration $T$ corresponds to $\\frac{T}{\\eta^2}$ steps of Adam in discrete time. A similar approximation theorem for the slow SDE of SGD can be found in [1].  \n\n**W1 part3: I wonder why there exists a factor $\\sqrt{\\eta}$ before the diffusion term of the conventional SDE in Line 129. **\n\n**A:** This is a good question. The factor $\\sqrt{\\eta}$ in front of the diffusion term is standard in SDE analyses of SGD. See [4,1,2,5,6]; they all add this $\\sqrt{\\eta}$ factor in their conventional SDEs.\n\nIn the SGD update, the gradient term is scaled by $\\eta$. Thus, the corresponding diffusion term in the SDE, which represents the stochastic noise, scales by $\\sqrt{\\eta}$. This scaling arises from Euler's first-order discretization of the stochastic process. Putting it mathematically, for an SDE \n\n$$dX_t = a(X_t)dt + b(X_t)dW_t,$$ \n\nits Euler discretization with a time step $\\Delta t=\\eta$ is\n\n$$X_{t+\\eta} = X_{t} + \\eta a(X_{t}) + b(X_{t}) (W_{t + \\eta}-W_{t}),$$\n\nwhere $(W_{t + \\eta}-W_{t}) \\sim \\mathcal{N}(0,\\eta\\mathbf{I})$. We want the diffusion term and drift term above to coincide with those of a single-step SGD update. For an SGD update, we write out its update rule at step $k$ as\n\n$$\\theta_{k+1} = \\theta_k - \\eta \\nabla \\mathcal{L}(\\theta_k) + \\eta z_k,$$\n\nwhere $z_k := \\ell_{k}(\\theta_k) - \\mathcal{L}(\\theta_k)$ and we define the covariance matrix $\\Sigma(\\theta_k):=\\text{Cov}(V_k,V_k\\mid\\theta_k)$. We have that $\\text{Cov}(\\eta V_k,\\eta V_k\\mid\\theta_k)=\\eta^2\\Sigma(\\theta_k)$. We identify $t$ with $k\\eta$, therefore, if we set\n\n$$a(\\theta_k) =  \\nabla \\mathcal{L}(\\theta_k), \\quad b(\\theta_k) = (\\eta\\Sigma(\\theta_k))^{1/2},$$\n\nthen the diffusion and drift terms of SDE will match those of the SGD update. Hence, this motivates the approximating equation\n$$dX_t = \\nabla \\mathcal{L}(X_t)dt + \\sqrt{\\eta}\\Sigma(X_t)^{1/2}dW_t.$$\n\nThus, the factor$\\sqrt{\\eta}$ is necessary to ensure that the diffusion term is correctly scaled and reflects the noise variance, maintaining consistency with the overall SGD dynamics. We hope this clears up your concerns.\n\n**W2 part1: Can you explain the connection to 'The implicit bias of Adam on separable data'?** \n\n**A:** Thank you for this good question. We have cited and discussed this work in our paper (2nd paragraph of Section 2). The paper you cited is, in fact, **orthogonal to our research**. Their work studies overparameterized linear logistic regression, where the minimizers are at infinity and Adam converges only in direction. There is no minimizer manifold in their setting. By contrast, our analysis focuses on settings with finite minimizers (e.g., logistic regression with L2 regularization). Therefore, their results do not apply to our setting, nor do ours apply to theirs. Our setup is arguably more realistic since real-world models (1) are non-linear and (2) have finite minimizers.\n\n**W2 part2: I suggest that the authors could design some experiments to directly support their main theorems on the trace of the Hessian matrix.**\n\n**A:** We would like to point out that our paper already provides experiments on matrix factorization to validate the implicit bias of Adam (Appendix B). In particular, we plot the curves of $\\text{tr}(\\mathbf{H})$ and $\\text{tr}(\\text{Diag}(\\mathbf{H})^{1/2})$ for Adam, and demonstrate that Adam is minimizing $\\text{tr}(\\text{Diag}(\\mathbf{H})^{1/2})$ rather than $\\text{tr}(\\mathbf{H})$ (Figure 4, top-right). This causes a negative effect on generalization compared with SGD (Figure 4, bottom-right).\n\n---\n\n**Q1: Explain more about ''Slow SDE'' for both SGD and Adam.** \n\nThanks for your feedback. We will break down your question into 3 sub-questions and answer them one by one. \n\n**Q1.1: Could the authors explain more about ''slow sde\" for both the SGD and Adam? For me, the conventional SDE is clear. However, the \"slow sde\" is relatively hard to follow.**\n\n**A:** In the following, we demonstrate the difference of slow SDE with conventional SDE from two aspects.\n\n- **The slow SDE traces the trajectory of SGD in a longer horizon compared with the conventional SDE**. As depicted in Appendix A, conventional SDEs can only track iteration for $O(\\eta^{-1})$ time, while slow SDEs can approximate the process after the parameter gets close to the minimizer manifold (up to $O(\\eta^{-2})$). This longer time horizon makes it much easier to see the long-term regularization effects of SGD directly from the SDE.\n- **Also, slow SDE traces the trajectory of the projection of the model parameter onto the manifold, rather than the parameter itself.** This allows us to better understand how the implicit regularization effect gradually influences the global minimizer that SGD centers around over time.\n\n**Q1.2: It's difficult to understand why the trajectories of SGD can be characterized in this form. What do the diffusion term and the drift term indicate?**\n\n**A:** We provide an interpretation of slow SDE beginning at line 168. Next, we first give a comprehensive explanation of the diffusion term, drift term, and slow SDE.\n\n- The **diffusion term** captures the stochastic noise from mini-batch gradients and other gradient noise (e.g., label noise). It injects random perturbations along directions into the tangent space of the manifold, driving the parameter to wander around the manifold instead of settling at a single point.\n- Meanwhile, **drift term** serves as a deterministic force that pushes the solution toward flatter, more stable regions on the manifold. When the label noise condition is satisfied, the indication of the drift term become clearer, where it can be interpreted as reducing a regularization term $(\\text{tr}(\\nabla^2 \\mathcal{L})$ for SGD, and $\\text{tr}(\\text{diag}(\\nabla^2 \\mathcal{L})^{1/2})$ for Adam) inside the manifold.\n\nThe mechanism behind slow SDE is quite complicated. In simple terms, in normal space, SGD behaves like an Ornstein-Uhlenbeck (OU) process locally, whose mixing time is $\\Theta(\\eta^{-1})$. After writing out Taylor's expansion of this process up to the third order, further analysis can show that the third term in this Taylor expansion creates a $\\Theta(\\eta^2)$ velocity in the tangent space, and there is no push-back force in the tangent space. Thus, the small velocity accumulates over time, making SGD continuously move on the manifold to a certain direction. And slow SDE characterizes this slow dynamics.\n\n**Q1.3: Why is the current updating direction of SGD (from my understanding, this should depend on its local information) related to its convergence point?**\n\n**A:** **Actually the convergence point defined by $\\Phi$ is local information.** This is because the iteration stays very close to the manifold during the timescope of slow SDE (refer to Figure 3\\(c\\) in Appendix A for an illustration). It is important to note that the slow SDE serves as an approximation of the SGD only when the parameter is near the minimizer manifold.\n\n**Q2: Explain how you deal with the second term in line 306 of the proof sketch?**\n\n**A:** Here we just open this term using the definition of  $\\partial F(\\theta)[\\mathbf{M}]$ in Section 3. Specifically, we have    \n\n$$\\nabla \\left(\\mathbf{S}\\right)\\left[\\nabla^2\\mathcal{L}\\right] = \\sum_{j,k}\\nabla\\left(\\mathbf{S}_{jk}\\right)\\cdot\\mathbf{H}_{jk},$$\n\nand then plug in the expression of $\\mathbf{S} = \\text{Diag}(H)^{-1/2} + O(\\epsilon)$ into this equation gives\n\n$$\\nabla \\left(\\mathbf{S}\\right)\\left[\\nabla^2\\mathcal{L}\\right] = \\sum_{j}\\mathbf{H}_{jj}^{-1/2}\\mathbf{H}_{jj} + O(\\epsilon \\text{tr}(\\mathbf{H})),$$\n\nand take $\\epsilon \\to 0$ gives the final results. The more detailed calculation can be found in Appendix G\n\n**Q3: Can we include first momentum in the slow SDE?**\n\n**A:** It depends on how close $\\beta_1$ is to 1.\n- As our paper assumed, If $b_1 = 1 - \\beta_1$ is a constant (independent of $\\eta$), the error stays $O(\\eta^{1/4})$ and we cannot add $\\beta_1$ into slow SDE. (Assumption C.6,line 646)\n- If $b_1 = o(1)$ (e.g. $O(\\eta)$), then $\\beta_1$ would change the slow SDE form, but this case is less practical, since $\\beta_1$ is typically not greater than $0.9$ in real-world areas. So we left it for future work."}]

Author rebuttal:

Meta-review: This paper investigates the implicit bias of Adam and demonstrates that the optimization trajectories of Adam, after reaching the minimizer manifold, can be approximated by a continuous SDE. The reviewer agrees that understanding the implicit bias of Adam is an important problem, and the paper presents solid results. A minor comment is that the requirement relating the parameters $\beta$ and $\eta$ may limit the generalizability of these findings to broader settings. The authors are encouraged to incorporate the reviewer’s constructive feedback and discuss potential limitations in the revised version.
Decision: Accept (poster)