Fast Equilibrium of SGD in Generic Situations
Authors: Zhiyuan Li
Yi Wang
Zhiren Wang
Venue: ICLR 2024 poster (2024)
Keywords: optimization
fast mixing
fast equilibrium
SDE
SGD
large deviation principle
Primary area: optimization

Abstract: Normalization layers are ubiquitous in deep learning, greatly accelerating optimization. However, they also introduce many unexpected phenomena during training, for example, the Fast Equilibrium conjecture proposed by (Li et al.,2020), which states that the scale-invariant normalized network, when trained by SGD with $\eta$ learning rate and $\lambda$ weight decay, mixes to an equilibrium in $\tilde{O}(1/\eta\lambda)$ steps, as opposed to classical $e^{O(\eta^{-1})}$ mixing time. Recent works by Wang & Wang (2022); Li et al. (2022c)  proved this conjecture under different sets of assumptions. This paper aims to answer the fast equilibrium conjecture in full generality by removing the non-generic assumptions of Wang & Wang (2022); Li et al. (2022c) that the minima are isolated, that the region near minima forms a unique basin, and that the set of minima is an analytic set. Our main technical contribution is to show that with probability close to 1,  in exponential time trajectories will not escape the attracting basin containing its initial position.

1 Published as a conference paper at ICLR 2024
FAST EQUILIBRIUM OF SGD IN GENERIC SITUATIONS
Zhiyuan Li
Yi Wang
Toyota Technological Institute at Chicago chiyuanli@ttic.co
Johns Hopkins University ywang2 jhu edu
Zhiren Wang" Pennsylvania State University zh i renwêpsu.edu
ABSTRACT
Normalization layers are ubiquitous in deep learning, greatly accelerating opti- mization. However, they also introduce many unexpected phenomena during training, for example, the Fast Equilibrium conjecture proposed by (Li et al. 2020), which states that the scale-invariant normalized network, when trained by SGD with 1j learning rate and weight decay, mixes an equilibrium in Ó nx
steps, as opposed to classical e eO((n)) ') mixing time. Recent works by Wang & WWnn (2022); Li et al. (2022c) proved this conjecture under different sets of assumptions. This paper aims to answer the fast equilibrium conjecture in full generality by removing the non-generic assumptions of Wang & Wang (2022): ; Li ct al. (2022c) that the minima are isolated, that the region near minima forms a unique basin, and that the set of minima is analytic set. Our main technical con- tribution is to show that with probability close to in exponential time trajectories will not escape the attracting basin containing their initial position.
INTRODUCTION
Normalization layers are ubiquitous and play a fundamental role in modern deep learning, e.g. Batch Normalization [offe Szegedy 2015), Group Normalization (Wu He 2018), Layer Nor- malization (Ba et al. 2016). and Weight Normalization (Salimans & Kingma 2016). Normalization layers not only greatly facilitate optimization and improve trainability, it also brings intriguing new optimization behaviors to neural networks. For example, & Arora (2020) showed that normalized networks can be trained by SGD with exponentially increasing learning rates, because training with exponentially increasing learning rates turns out to be equivalent to training with constant learning rates but with weight decay turned on, as shown in (2). Here Vy E Rdd is the parameter of a neural network after the k-th step and is updated by
Øb+1 + (1 A)T V/ LB, (Ix),
where À and 1 are respectively the weight decay parameter and the learning rate, and LB, is the loss function evaluated using randomly chosen mini-batch B3.-
The result of Li Arora (2000) holds not only for normalized networks but more broadly for all scale invariant training losses, which is popular abstraction of normalized networks in optimization analysis. Mathematically, scale invariance refers to the following property of the loss: Lg(er) =
[g) Vc 0, Rdd and every batch B.
Later, et al. 20200; Wan et al. 2021 discovered that it is the intrinsic learning rate yjλ that
controis the long-term convergence behavior for SGD scale invariant loss with weight decay, (I)
The approach that Li et al. (2020) takes to study (T) is to approximate by the stochastic differential
equation (SDE) model Li et (017 2019). which is quite common in literature.
dx, (-nVL(X)) m)XX))dt (X,)
"The authors are listed alphabetically.

---

steps of SGD Published as a conference paper at ICLR 2024
Here is the average K Ek=1 Ls, over all random batches, o= vR (VLB. VL) k=1 is K matrix, and B is the K K-dimeenional Wiener process. Scale invariance of Lø implies that L is ccale-invariant and that o is (-1))hommgeneouse i..
M(er) L ((x), c c= Ve 0
Li et al. (020) further proposed the following Fast Equilibrium Conjecture for the SDE approxima- tion of SGD.
Conjecture 1.1. [Fast Equilibrium Conjecture] (Li et al. 2020 F(X, input) denotes the output of a neural network NN with parameters X, and Xt denotes the value of SDE (2) at time starting from initial parameter Xu. Suppose NN has normalization steps so that the F F(X, imput) is scale- invariant in X, i.e. F(X, imput) = F(cX, input) for all O. Then for all input values input, the probability distribution of F(X,, input) stabilizes to equilibrium in O 7x updates.
Experiments where the empirically observed rates of convergence are polynomial were contained in the original paper Li al. (2020) where the Fast Equilibrium Conjecture was first asked. The rate O is considered to be fast because according to Langevin dynamics, the time it takes to converge
to the Gibbs equilibrium of exponential order ((((A)) b), This can be done by following a similar analysis to those in (Bovier et al. 2004,, (Shi et al.) 2020p. The works by (Bovier et al. 2004) and Shh et al. 2020) dealt with models without normalization, and the convergence times there are of order e ((nA)) Using the similar method as in (oviier et al. 2004) and (Shi et al. 2020), when
normalization is used the convergence time can be shown to be of order (((A)) 2) This is because Li al. 2020) proved that the intrinsic learning rate ε/λ is replaced by an effective learning rate
(Y's in Li et al. )020 for the renormalized parameter vector, which is of order ((()))))
The recent paper (Li al. 2022c), using a mathematical framework from (Li et al., 2022b), estab- lished the fast equilibrium conjecture for 7/ - 0 under mixed set of generic and non-generic assumptions. See (Damian et al. 2021), (Gu et al. 2022) for more work on analyzing the dynamics of SGD near the manifold of minimizers. The goal of the current paper is to remove the non-generic ones and thus provide a general proof in the aforementioned range of parameters.
1.1 NOTATIONS AND ASSUMPTIONS
To introduce assumptions from previous authors as well our results, we need set up few notations first. Let r C Rd {0} be the set of local minima of L. Notice that by (3) I is cone, i.e. x T if and only if cx E for all c 0. For all r > write r, far E [x| r). In particular I, is subset of the unit sphere gd- {|æ] = 1}.
In general. I may have multiple connected components. Decompose = r where each I' is connected cone. We then write r = L [|r] r). Then l are the connected components of T In particular, there are only finitely many I and we index them by i 1, ..... m.
In addition to the scaling properties (3) guaranteed by the use of normalization, (Li et al. 2022ch made certain assumptions, which we will need the following.
Assumption 1.2. The functions I. and satisfy:
(i). (Scale invariance) The scaling rules in (3) hold.
(i.. (Regular critical locus) Each loss function LB, is C* IR (0), the critical points of form a C2 submanifold Q. For all x E Qt L( of rank dim T,Ω.
(iii). (Controllabilitt) For all Ií, (24 lim4-po X,, with X, being the solution to the deterministic gradient descent X, = LVL(X,) with initial value x.
T.I,, Here and below, 4(x)
By Arora et al. 2022 Lemma B.15), under Assumption 1.2 & (ii), is well defined and C differentiable on neighborhood of I' as long as L is 4 differentiable. We also note that in general
Different connected components of Q are required have the dimension.
2

---

Published as a conference paper at ICLR 2024
the noise structure does affect the convergence rate. But as long as Assumption 1.2 (iii) is satisfied, the noise structure won't t affect the asymptotic order of the convergence rate.
On the other hand, we will not need the following assumptions.
Assumption 1.3. I satisfies:
(i) (Unique basin) 11 is compaet and connected;
(i.). (Analyticity) I' is a real analytic manifold and Tr E is real analytic function on Ro where Σ
Restricting to an attracting basin U of I and assuming both Assumptions 1.2 and (2022c) proved Conjecture 1.1 when λrj - 0 in the natural range of 1/ < O()) 0 and the parameter Xo is initialized within U. Note that since U is attracting basin, St and I coincide in U and thus Assumption [.2 is equivalent to (Li et al. 2022c Assumption 2.1) for the purpose of that paper. Note that I' is always a submanifold of S2.
et al.
Remark 1.4. All three assumptions in Assumption
1.2 very natural for the following reasons:
As remarked earlier, the acale-invariance 3 is consequence of the use of normalization steps inside neural networks.
It is a widely used assumption, at least in the case of local minimizers, that the locus is a manifold for overparametrized neural networks, for example in (Fehrman al. 2020; Arora al. 2022 Li
el al. (2022b). For the locus of global minimizers, this assumption was proved by Cooper 2021). As remarked in (Li al. 2022b Cooper [02T], a main reason for the local minimizers to form manifolds is the overparumelrizatin of modern neural networks. Şimsek et al. (2027) further iden- tifies the reason as symmetries arising from overparametrizatio. fact, they studied loci of critical points that are not necessarily minima and proved that symmetry- conmetry-induced critical points form manifold that satisfies Assumption 1.2 (ii).
The philosophy behind Assumption 1.2 (ii)) is that the generation of random batches in training is independent of the aforementioned symmetries in the setup of the neural network, and thus generi- cally should not live in subspaces that are invariant under such symmetries. In particular, the same symmetries are generically capable of move the noises from the random batches to spun a tangent space of the same dimension that of the local manifold of critical points.
Remark 1.5. On the other hand, both conditions in Assumption 1.3 are non-generic:
Certain evidences from (Drxdle et al. 2018) suggests that all local minima appearing in realis- tic training generically come from connected relatively flat region of small variation in height, so empirically Assumption (1) could be a reasonable approximate assumption. However, the exper- iments in Draxler et al.] 2018 Fig. 5) shows at the same time that in many settings, this region not completely flat and contains non-trivial saddle points. particular, there could be multiple disconnected basins. In light of these, it is more reasonable work in the absence of assumption Assumption (i).
The analyticity of the I and Tr E depends that of the activation functions chosen in the neural network. While many popular activation functions are real analytic, one may always choose to use functions that are differentiable but not analytic, in which case Assumption 1.3 (ii) is in general not guaranteed.
the same natural range as in (Li et al. 2022c), assuming only the generic conditions from Assumption 1.2 In particular, we will introduce two arguments that respectively remove both hypothesis (Unique basin and Analyticity)
In this paper, we will give general proof of Conjecture I.1
in Assumption I.3
Here we would like to provide comments on why Assumption 1.3 are restrictive. Assuming ana- lyticity is restrictive because the regularity of the loss function is decided by that of the activation function. Even though popular activation functions such as Sigmoid are analytic, a priori one could smooth but not analytic functions. The one basin assumption is restrictive empirical evidence of proof that L only has one basin. In fact, the experiments at the end of the paper suggests that there are multiple basins.
do not
3

---

Published as a conference paper at ICLR 2024
We would also like give remarks on why the three assumptions in Assumption 1.2 essential. (i) is essential because without this assumption, the SDE would not be equivalent to a SDE on the sphere $ which is crucial to analysis. Without this assumption, similar analysis can probably be formulated on Rd instead of the $ coordinate but there will be new technical obstacles to overcome. Since the original fast equilibrium conjecture was asked for normalized neural nets, restrict our study to the current setting. (ii) is important because if not a trajectory may stay near critical point (for example saddle point) for a very long period of time, would not be able to converge within a polynomial time. Finally, the reason why we need (iii) is that if the span is the whole tangent space, but instead a subspace of the tangent space, then the diffusion will be restrained to this subspace, which a priori may be very fractal and existing mathematical theory is not sufficient to guarantee unique equilibrium limit.
1.2 MULTIPLE EQUILIBRIA WHEN BASIN IS NOT UNIQUE
It is worthy to explain in more details what happens when the basin is not unique, i.e I' has multiple connected components and Assumption 1.3 (i) fails. In this situation, our analysis generalizes the work of Wang & Wang (2022) and reveals three-stage equilibrium phenomenon. The most im- portant property of this phenomenon is the mismatching between practical training and theoretical bounds: the equilibrium distribution of network parameters observed in the time window under realistic budget is both local in space and temporary in time. It is concentrated near the bottom of the same attracting basin containing the initial parameter, and differs from the eventual global Gibbs equilibrium that the distribution of parameters will eventually converge to in exponentially long time. This phenomenon interprets the gap between the empirically based Conjecture 1.1 and the previous theoretical estimate from e.g. (Bovier et al.) [204]] Shi et al.) [020)). See (Frankle et al.) 2020), (Gupta et al., 2019) for more work about the iterates stay in the same basin for a significant amount of time when starting from the same initialization.
The major short-come of (Wang& Wang 2022) is that, while not relying on the uniqueness of the basin, the arguments therein are subject to other non-generic assumptions, namely: (1) all basins are isolated points; (2) the noise a is a standard isotropic Gaussian noise.
Our methods allow to remove these assumption simultaneously together with Assumption 1.3 This made possible by avoiding using the cami-ccassccal analysis of spectra of differential operators, which was developed by Simon (Simon) 1983) and used in an essential way by previous authors in (Bovier et al. 2004]; [hh et al. 2020 Wang & Wang 2022). Instead, our method is purely proba- bilistic and predicts that the exiting time from a given basin is exponentially long. This method is based on an adaptation of the large deviation principle of Dembo & Zeitouni (2010).
2
STATEMENT OF MAIN RESULTS
2.1
PRELIMINARIES ON SDE MODEL
A polar coordinate system has been adopted in (Li al. 2020) to study the SDE model 2 For this
purpose, denote by X, [X+] the unit renormalization of X,, and |X,1
Xx
By Li et al.
2020 Theorem 5.1), 22 is equivalent
1 -1
dx, Y x(x
2 76 Tr TrE I(X.
(4)
d% 4n
= + +2 E(T
(5)
di
Recall that = aa
T is a d d positive semidefinite symmetric matrix.
One may view the motion of X, , as an intrinsic one inside the unit sphere Sd-1 instead of one inside Rdd. From this perspective, (Wang & Wang 2022. Theorem 3.1) shows that can be rewritten as an intrinsic SDE on Sd-1
dX, - (VL(X)dtt ( (X.)dB
(6)
where 6-)² is a tensor field along sd-1 whose value is given by the restriction of o()) and V is the gradient operator on gd-1, (See the remark after WWag & Wang 2022 Theorem 3.1) for the
4

---

Published as a conference paper at ICLR 2024
meaning of being intrinsic. In particular, df, a (X are viewed vector fields along 1)
Remark 2.1. Instead of the term o (Xx in (4) and O the papers (Li et al. 2020 Wang &
Wang 2022) actually used the restriction of x to Sa-1, However, these two expressions are equivalent as Wiener processes because Σ = 00
Since (6)
perturbation with Brownian noise of the gradient flow dX, -y - VI.(x) with varying learning rate % makes sense to first understand the constant speed gradient flow
X (I)
(7)
which is an ODE on the compact manifold gd-1, Following earlier notation, the local minima of L on gd-1 is F1 and has connected components ri.
Write Uf g for the attracting basin of Tí, i.e. the set of Xo Sd-1 such that the solution X to ( with initial value X, satisfies lime-+oo X, ri).
Lemma 2.2. Under Assumption (i),, the U| for 1, 1.... m. are disjoint open sets of Sd-1, Moreover their union L Uj has full volume in gd-1 and 1d-1 Ui is proper submanifold of
The proof of Lemma 2.2 standard and we left to the reader. The key observation is that the complement 3d-1 L Uj is the union of attracting basins of the connected components of critical points that attracting basins are proper submanifolds.
local minima. By Assumption 1.2 (i), those critical points are saddle like and their
Note that L is constant on I and 1", and L(T)) L(T)) is the minimum of L inside Ui. For > 0, write
U = e U L(x) < L(U1) + +
Define cones U fx E [0) e Ui}, U (s E Rd (0) 1 E U...). Then U attracting basin of T under the gradient flow
X, =- VXX))
(8)
By Arora al. 2022 Lemma B.15), under Assumption 1.2 (i), U is open and the function ((x) lim + X₁ is differrentiabe on U.
2.2 MAIN RESULT
Definition 2.3. We define the Lipschitz distance between two probability measures μ1, metric
space X as
If f pdde),
distráp(uu, v) sup $1
where the Lipschitz norm of a function &p is given by I|<l| == max (Ieelloo,supp supazy
(=(z))p(1) d(z,z)
We are now able to state our main theorem, which a mutual reinforcement to both (Li et al. 2022c Theorem 5.5) and (Wang& & Wang 2022] Theorem 4.6).
Theorem 2.4. Under Assumption 1.2 for all > and compact interval [p= P+ C (0, 20), there
exist constant >> O and set A
SUm Ui C Sd-1 of volume volga ( (A)
i
such that for
all & O()) & (1), the following holds:
For all initial parameter
A. all growth rates K such that Rd with |ool E [p_, and spl
K + 20 as yd + O, and all time values
K log(11 + 2)
t tE
e vix
the random trajectory with initial value Do satisfies dist (Pxx=z((X)) v²) < e,
where i is a probability measure supported on the attractor I of the unique attracting basin containing To, and v² only depend on L,, and i.
5

---

Published as a conference paper at ICLR 2024
Here volga-1 1 is the renormalized volume the sphere g so that the total mass of Sd-1 is 1.
3 REMOVAL OF ANALYTICITY ASSUMPTION
In this part, we will prove that Assumption 13 (ii) on analyticity (which is (Li et al. 2022c Assump-
tion 5.3)) is unnecessary for (Li et al. 2022c Theorem 5.4), and thus (Li al. 2022c Theorem 1.2)
holds without such an assumption as well.
Proposition 3.1. Under Assumption 1.2 and Assumption 1.3 (1), the conclusion of (Li et al. 2022c Theorem 1.2) hold.
For this purpose, we temporarily adopt the setting (Li et al. 2022c) for now. other words, Assumption 1.3 (i) is being assumed and [3 is compact connected submanifold of g and I = frz : >> 0; r E }. By eL al. 2022c Theorem 4.1) it suffices to prove the (qualitative) mixing of the SDE Using the notations from (Li et al. 2022c Chapter E.3), this SDE writes:
2022c Equation (13)) I towards unique invariant probability measure.
dY, = fo(Y))
Σ fx())) -dBKr.
where fA are certain vector fields along the radial cone r S ( and B,,,, 1, K are iid. Wiener processes. Indeed, -,,-.. S are orthogonal to the radial direction with
fz) = (z))) 1 &
and they span Tl'l = Tx for all æ E r, where r, {x e e:: |x| = and fo has the form
K
1 ) - 0 [
fo(x)
(f(())((()))
i=k
In the proof' from (Li et al. 2022c analyticity is only used in Chapter F.4, when Tr E is not constant on F1. In this case, it was proved there (without using analyticity) that I, == { r r- < | < ]+} with
min (Tr Σ) rr max Σ)+
(10)
the unique invariant control set for the control problem corresponding (
Instead of using Kliemann's 's condition (Kliemann 1987). which requires the vector field Lie alge- bra I generated by fo. ... fN be of maximal dimension at all points in I, and the reason for the need of analyticity in (Li et al., (2022c), we will use Kmold-Kliemane Kliemann 1987), which only requires that the vector field Lie algebra to be of maximal dimen- sional at point in r.. This condition is true because the projection of fo to the radial direction is not constantly 0 Tr Σ is not a constant. (Otherwise I', wouldn't the unique invariant set.) Under this condition, (Arnold & Kliemann, 1987 Theorem 5.1) proved that there is unique in- variant probability measure / supported on r.. The measure v then has to be ergodic. Moreover, (Amold & Kliemann| 1987] Theorem 5.2) showed that is absolutely continuous with respect to the Riemannian volume on the manifold l' In particular, v(r..) ) = 0 and v(T.) = 1.
condition (Amold &
The main issue is to prove the convergence of the distribution St, towards denotes the measure of all trajectories of solutions to (9) at time / starting from x E T,, A priori, such convergence is only known to hold for 7 Jo" St by ergodic theorem.
L - oc, where St.z
Applying Duflo & Revuz 1969 Theorem II.4) to (T,, it suffices to check two conditions to guarantee St,= v (in total variation distance):
(Harris's recurrence condition) For all sets A with (AA) 0 and all E T.,
Fxo=x 11(x))) = =00
(11)
(Regularity condition) The absolutely continuous component of Sr, respect to called 89 satisfies limp-zoo 80 (T.) for all x E T..
f,

---

Published as a conference paper at ICLR 2024
Let us first verify (Harris's recurrence condition). Define D Px by D {v T. fo(a)) $ [=]]-
Then open in T.. For all æ r. and T., we define the random variable Ta,A 2 0 to be the first entering time into A for trajectory of (9) starting at Lemma 3.2. P(Ta,D < oo) for all T..
Proof of (Harris's recurrence condition). By (Arold & Kliemann, 1987 Theorem 6.1), (11 holds for all E D n intl', = D. By Lemma 3.2 it then holds for all E P.. This verifies Harris's recurrence condition.
We now verify the (Regularity condition): In addition to the Lie algebra and set D, define Lie algebra lo S TT and an open set Do by
lo = span(fi, ... fN) + [r,4, Do = {u E T, dim lo(z) = dim r.). It is easy to see that lo is indeed a Lie algebra and Do is relatively open subset in I'. Lemma 3.3. The set Do has non-empty interior.
We postpone the proofs of Lemma 3.2 and 3.3
to Appendix
Lemma 3.4. For all x Do and
S
(T.)
>
Proof. By (Ichihara & Kunita 1974 Lemma 1), every x E int the second order differential operator on the right hand side of (9) elliptic (non-degenerate) at æ. The lemma follows.
Proof of (Regularity condition). By (Harris's recurrence condition) and Lemma 3.3 for all T., P(Tz.Do < oo) > 0, and thus there exists fo(x)) > 0 such that IP(T,,Do < to(x)) >> 0. In other words, on a subset S2,.D. C Q of stochastic incidences with P(S2g.Do)
0, there exists
TTD,, 2 D. (w) E (0,4((z) such that the solution starting at Y Y())
satisfies Y (T',D,) Do- (T.) > 0. This shows that the
By Lemma 3.4 for all / >> (o()), St, (T.) 2 Sa So Si St-T..D y P.
statement For v-ae. x, Sh (T.) for all 0" false. By (Duflo & Revuz 1969 Proposition, p235), this guarantees (Regularity condition).
We have by now completed the proof of the mixing property S.,z - v under the generic Assumption 1.2 and the Assumption T.3 (i).
4 REMOVAL OF UNIQUE BASIN ASSUMPTION
We now stop assuming Assumption 1.3 (i) and decompose r L where each I is connected
cone. Write r
- (1al E [R_, R+]3 where we fix R_ <
(n /mmn -11 E(x) and R4 >
2A
1) max 111-1 TV To) near these bounds.
2A
Note that I' f|o| E
[[ [22mi Tr x2(r))
(2()))
an invariant control set of (2
2x
Moreover, was proved in (L al. 2020) that for a given initial radius [zo], the radius |X,] of (2) starting at will be almost surely isside [R_, R.] for 2 ((nA)=1)
Write r 1 [|æ] r} and r I Then r are the connected components of the manifold 1, In particular, there are only finitely many I,'s. Write dim I For s > 0, write Ui, for the p-eeighborhood of r in ri, U.. rUí,p and Ut Urein_.n,) U
Fix from now on a sufficiently small parameter Po such that Ui. are disjoint for distinct i's, and ri is the set of all critical points of L inside Ui, U0 This is possible because of Assumption 1.2 (ii). Proposition 4.1. For all sufficiently small pı (to be determined later), given K 1, 6 > there and so. lol E AK,e with probability > - e, the trajectory of (2) starting at Zo will remain inside U, U.. for E [0, Cacs(1)))1) some constant Caes-
exists subset A AR. S gd-1 with mia 1 ((Ag,e) 1 such that for all no with lool E 1 K]

---

Published as a conference paper at ICLR 2024
The proof of the proposition, which omit, is simple combination of (Li et al. 2020 Equation (7)) and the proof (Wang & Wang 2022 Theorem 4.5).
The proposition allows us assume that our initial point is inside one of finitely many basins U*.pp To prove the main result, it now suffices to make two observations: First, with very high probability, the trajectory will not escape from the basin in exponential time. Second, as long as the trajectory remains in the basin, its distribution always mixes towards unique probability measure v¹ supported at the bottom I the basin.
The first property is stated Proposition 4.2 below.
Proposition 4.2. There exist C > O, and sufficiently small po > D1 > O, such that for all i and U in the regime 1) (()) S O(1), the solutions X, to with initial position xo satisfy
F5 #.p1'
for t E e(nA) 2D =1
lim P Pxo=ao (X, remains in U t.ro X0-+0
The convergence is uniform with respect the inital data xn-
It is similar in nature to (Wang & Wang 2022 Lemma E.6) but requires a more sophisticated proof. This is the main theoretical component this paper. The argument will be based on the large devia- tion principle of Dembo-Zeite Zeitouni in Dembo & Zeitouni (2010) Chapter 5), which was an adaptation of (Freidlin & Wentzell 2012, Chapter 6). The reason for which Freidlin-Wentzell
original theory cannot be applied here like in (Wang & Wang 2022) is that the diffusion in the SDE system (44), 45 is degenerate. Dembo-Zeitouni work allows degenerate diffusions. However, further modifi- cations to (Dembo & Zeitouni 2010) needed in case as the first order drift in (45) depends on the " -2 We will treat 2/ as a control variable. The full proof will be in Appendix B
The second property follows from the main results (Theorem 5.1 & Theorem 6.7) in (Li et al. 2022ch) and is restated as Proposition 4.3 here with an additional emphasis on uniformity. A more detailed discussion
be found in Appendix C Recall that (Li al. (022ch) also assumes As- sumption 1.3 (ii), but that can be dropped by the discussion in Chapter above.
Proposition 4.3. Under Assumption 7.2 For all > and sufficiently small po > P1 > O, such that in the regime n S O(x) S 0(1) and n, - O, the following holds: For each index i and for all initial parameter xo E U #mm the distribution of all trajectories
{e)
K 1o d of 2
Jogis
2d
that do not leave U' converges in distribution to the trajectories {X.} to fixed SDE model (the Katzenberger model) supported on I' with initial position ((zo). The convergence is uniform in zo- Moreover, as K - oo the trajectories {Xi} are uniformly mixing (with respect to different an 's) towards a fixed equilibrium measure v².
Our main theorem, Theorem 2.4. then follows from combining Propositions
and 4.3
Proof of Theorem 2.4 By Proposition 4.1 after ignoring ( time the beginning, well
o(1) portion of stochastic incidences. One may assume Do E U #.p for some i. For within the range in the statement of Theorem 2.4| again after ignoring an o(1) portion of incidences, one may assume that all trajectories under consideration stay within U/,,, up time t. Because of the lower bound for
I, one may consider window of length {o that ends at / where K + xo. The distribution of trajectories along this window is the average of distributions over different initial positions. By Proposition 4.3, all such components uniformly mix toward v².. The proof is completed.
[
EXPERIMENTS
Mixing on local manifold: The key technical observation of this paper (Proposition 4.2) is that the distribution of trajectories with an initial position is trapped locally in the attracting basin containing the initial position during any practical observation windows. Using the method from (Wang & Wang 2022 Fig. 13), this observation is supported by the experiment below: using reduced MNIST dataset with only 1280 samples and a small CNN with 1786 parameters (so that the model
still overparametrized), we ran 15 independent instances of SGD, À " 32 for each of two
8

---

Published as a conference paper at ICLR 2024
randomly chosen initial parametrizations, Each instance lasts 0.8 million steps of SGD. A smilar experiment was ran for reduced CIFARIO dataset with 1280 samples, a CNN model with 2658 parameters, 17 1024 À 32 and 1.28 million SGD steps. In order to show that the distribution arising from each initial position does stabilize toward an equilibrium and the two equilibria are different, track the variance within each group, and compare them with the average distance square over all pairs of point from different groups. Namely, denoting by {XE.) the i-th trajectory starting at initial point Ih where 1,2, we compute the following quantities:
(L) E X x),l? V(2(()= Eiys X2, X3,d12 V12(1) E X3.e..
Figures| 1 2 shows that V11 and V22 stabilizes near similar but different values, but V12 stabilizes a much bigger value. This suggests that the distributions of trajectories with starting point æ and 22 mixes towards equilibria whose support have similar scales, but these equilibria are far apart from each other.
100
3
62
Figure I: small MNIST: Variance comparison be- tween distributions with different initial positions
Figure 2: small CIFAR: Variance comparison be- tween distributions with different initial positions
02
50
Figure 3: small MNIST: Comparison between
Figure 4: small CIFAR: Comparison between
training losses before and after SWAP
training losses before and after SWAP
Prediction on the failure of stochastic weight averaging in parallel (SWAP): Our theory pre- dicts that if the local minima manifold of minimizes 1" has non-trivial geometry, that is, the average of parameters the manifold may fall off the manifold, then it might fail to decrease loss, or even increase the loss, once the SGD mixes to the local manifold.
We apply stochastic weight average over trajectories (SWAP) to the neural network parameters each step over the 15 independent instances with the same imitial position and compute the loss function the averaged parameter. SWAP, variant of stochastic weight average (SWA) from Izmailov et al. (018)), was proposed by Gupta et al. (020).. Figures 3 ssow that although the loss of the SWAP parameter improves the average loss over the independent instances at the beginning, the improvement quickly breaks after a couple thousands of training steps. This phenomenon verifies our theoretical prediction and also suggests that the support of the equilibrium is not a convex set but rather manifold of curved shape.
CONCLUSION
We give rigorous proof of fast equilibrium conjecture generie situations, removing previous as- sumptions that there is only one basin and the set of minima is analytic. The main technical con- tribution that we justify most the trajectories of SDE would not escape from one basin within exponential time. Instead of using spectral analysis, we adopt the large deviation principle type of argument. Possible interesting direction may include understanding the dependence of mixing time on dimension, architecture and noise structure.

---

Published as a conference paper at ICLR 2024
7 ACKNOWLEDGEMENT
Y.W. and Z.W. acknowledge respectively supports from NSF.
REFERENCES
Ludwig Arnold and Wolfgang Kliemann. On unique ergodicity for degenerate diffusions. Stochas- lics, 21(1):41-61, 1987.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent the edge of stability in deep learning. In Proceedings of the 39th International Conference on Machine Learning. pp. 1024, 2022.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint vrXiv:1607.06450 2016.
Anton Bovier, Michael Eckhoff, Véronique Gayrard, and Markus Klein. Metastability in reversible diffusion processes Sharp asymptoties for capacities and exit times. Journal of the European Mathematical Society, 006(4):399 424, 2004.
Yaim Cooper. Global minima overparameterizee neural networks. SIAM Journal on Mathematics of Data Science, 3(2):676-691, 2021.
Berfin Şimşek, François Ged, Arthur Jacot, Francesco Spadaro, Clement Hongler, Wulfram Gerst- ner, and Johanni Brea. Geometry of the loss landscape in overparameterized neural networks: Symmetries and invariances. Proceedings of the 38th International Conference on Machine Learming. pp.
2021.
Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers. In Advances in Neural Information Processing Systems, pp. 24449,27461, September 2021.
Amir Dembo and Ofer Zeitouni. Large deviations techniques and applications. volume 38 of Stochastic Modelling and Applied Probability. Springer- Verlag, Berlin, 2010. ISBN 978-3-642- 03310-0. doi: 10.10077978---644--000
Corrected reprint of the second (1998) edition.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred A. Hamprecht. Essentially bar- riers in neural network energy landscape. In Proceedings of the 35th International Conference on Machine Learning., volume 80, pp. 1308-1317, 2018.
M. Duflo and D. Revuz. Propriétés asymptotiques des probabilités de transition des processus de Markov récurrents. Ann. Inst. H. Poincaré Sect. B (N.S.), 5:233-244, 1969.
K.J. Falconer. Differentiation of the limit mapping in a dynamical system. Journal of the London Mathematical Society, 52-27(2)335633777 -322, 1983.
Benjamin Fehrman, Benjamin Gess, and Arnulf Jentzen. Convergence rates for the stochastic gradi- descent method for non-convex objective functions. Journal of Machine Learning Research, 21(136):1-44 2020.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode con- nectivity and the lottery ticket hypothesis. In In Proceedings of the 37th International Conference on Machine Learning, pp. 3259-3269, November 2020.
Mark I. Freidlin and Alexander D. Wentzell. Random perturbations of dynamical systems, vol- ume 260 of Grundlehren der Mathematischen Wissenschaften [Fundamental Principles of Math- ematical Sciences]. Springer, Heidelberg, third edition, 2012. ISBN 10.1007/978-
doi: Translated from the 1979 Russian original by Joseph Szücs.
Xinran Gu, Kaifeng Lyu, Longbo Huang. and Sanjeev Arora. Why (and when) does local sgd gener- alize better than sgd? In The International Conference on Learning Representations, September 2022.
10

---

Published as a conference paper at ICLR 2024
Vipul Gupta, Santiago Akle Serrano, and DeCoste DeCoste. Stochastic weight averaging in par- allel: Large-batch training that generalizes well. In The International Conference on Learning Representations, September 2019.
Vipul Gupta, Santiago Akle Serrano, and Dennis DeCoste. Stochastic weight averaging in parallel: Large-batch training that generalizes well. In 8th International Conference on Learning Repre- sentations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL
htttps://openrevil
net/ Foeee//oouuiiiirrgPPAA...
Kanji Ichihara and Hiroshi Kunita. A classification of the second order degenerate elliptic operators and its probabilistic characterization. Z. Wahrsshhiinlickkeiissheerrr und Verw. Gebiete, 30:235- 254, 1974. doi: 10. 10.1007//FF0033347
Sergey loffe and Christian Szegedy. Batch normalization: accelerating deep network training by reducing internal covariate shift. In Proceedings of the 32nd International Conference on Inter- national Conference on Machine Learning, pp. 448-456, 2015.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. In Ricardo Silva, Amir Globerson, and Amir Globerson (eds.), 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, 34th Conference on Uncertainty Artificial Intelligence 2018, UAI 2018, pp- 876-885. Association For Uncertainty in Artificial Intelligence (AUAI), January 2018. 34th Con- ference on Uncertainty Artificial Intelligence 2018, UAI 2018; Conference date: 06-08-2018 Through 10-08-2018.
Gary Shon Katzenberger. Solutions of a stochastic differential equation forced onto manifold by large drift. The Annals of Probability, pp. 1587-1628, 1991.
Wolfgang Kliemann. Recurrence and invariant measures for degenerate diffusions. Ann. Probab., 15(2):690-707, 1987.
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradient algorithms. In Proceedings of the 34th International Conference on Machine Learning. pp. 2101-2110, 2017.
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research, 20(40);
Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. In Interna- tional Conference on Learning Representations, 2020.
Zhiyuan Li, Kaifeng Lyu, and S. Arora. Reconciling modern deep learning with traditional optimiza- tion analyses: The intrinsic learning rate. Advances in Neural Information Processing Systems 33 preceeoceedsngs (NeurIPS 2020), 2020.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? - mathematical framework. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022a.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022b.
Zhiyuan Li, Tianhao Wang, and Dingli Yu. Fast mixing of stochastic gradient descent with normal- ization and weight decay. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022c.
Tim Salimans and Diederik P Kingma. Weight normalization: A simple reparameterization to accel- erate training of deep neural networks. In Advances in Neural Information Processing Systems, pp. 901-909, 2016.
Bin Shi, Weijie Su, and Michael Jordan. On learning rates and schrödinger operators. arXiv preprint 2:X4::2004.66977, 2020.
11

---

Published as a conference paper at ICLR 2024
Barry Simon. Semiclassical analysis low lying eigenvalues. I. Nondegenerate minima: asymptotic expansions. Ann. Inst. H. Poincaré Sect. A (N.S.), 38(3):295-33) -000,, 1983. ISSN 0246-0211.
Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamies: Learning dynamics of neural network with normalization, weight decay, and sgd. In Advances in Neural Information Processing Systems, volume 34, 2021.
Yi Wang and Zhiren Wang. Three-stage evolution and fast equilibrium for SGD with non-degenerate critical points. In Proceedings of the 39th International Conference on Machine Learning, pp- 23092-23113, 2022.
Yuxin Wu and Kaiming He. Group normalization. In European Conference on Computer Vision, 2018.
A PROOFS FOR SECTION 3
Proof. of Lemma 3.2 Instead of 9 define
dY, (Y,)) dt ( dBe,,
(12)
where (x) is the projection of fo(x) to T2T
For all 0 t D. fo fo at Y,, thus the solutions Y and Y to and 12) starting coincide up to D. Morcover, the trajectories of (22) stay in Fiz]. So Tx.D D is also the first entering time into DnI jr for (12. We claim that Dnr [r) non-empty subset of I |=-- Indeed, if this is not
true, then the trajectories Y and y coincide forever and stay inside |1 making I [la| an invariant control set which contradicts to the uniqueness of r.. Moreover, D D'l' [w] is relatively open in the
submanifold |z]. By Assumption (iii), T22 has non-degenerate diffusion on the submanifold
|a = [r|1'1 Thus Y enters D almost surely. This proves the lemma.
Proof. of Lemma 3.4 The proof is similar to that of (Li et al. 2022c Lemma F.18). Recall that Tr Σ is (-2) (-2)-hhmogennouus and assumed to be non-constant on 1. Therefore, there is open set
Vi E I2 and vector field f taking values in span ... fNN such that (V Tr Σ, f*) + on fo(x)) +
V.. By homogeneity, this is also true on rVi C Т, for all > 0. As Tr [(x)
3) and (x) = 0, this implies (z, fo(x)), f*) * 0 on rV.. By (Li et al. 2022c Lemma F.16), (z. [fo, f*]) * 0 on rV.. Note that [fo. f' E lo. and span(fi.-... .NN TT, on rVi by controllability. It follows that {rV1 : E (r-, C Do. which is sufficient to conclude.
B EXITING TIME WITH DEGENERATE DIFFUSION AND EXTRA CONTROL VARIABLE
In this appendix, we establish a probabilistic lower bound to the exiting time of a stochastic process from a basin (Theorem B.2IJ, based on a one-sided large deviation principle (Theorem B.12). Our proofs adapt those from the work of Dembo-Zeitouni (Dembo & Zeitouni 2010 Chapter 5) to more general setting. The main differences in the setting are:
1. The stochastic process in the basin is now governed not only by the current location and an Brownian motion, but also by extra control variable stated in equation (17):
2. The local minima set in the basin is no longer assumed to be unique isolated fixed point.
The diffusion in the stochastic process is allowed to be degenerate, which was the main novelty in the Dembo-Zeitouni theory compared to the earlier work Freidlin-Wentzell (Freidlin & WWetzell] 2012 Chapter 6).
This appendix will solely consist of mathematical analysis. With the exception of B.4 all notations chosen independently from those used in other parts of the current paper.
12

---

Published as a conference paper at ICLR 2024
B.1 PROPERTIES OF UPPER LARGE DEVIATION PRINCIPLE
In this section we define the notion of upper large deviation principle (upper LDP). This is the upper bound part of the large deviation principle defined in (Dembo & Zeitouni 2010 Chapter 1.2). The principles proved in (Dembo & Zeitouni 2010, Chapter 4.1 & 4.2), which allows to pass the LDP property between random processes, are still valid for upper LDP because the upper and lower bounds are treated separately in their proofs. The purpose of this section is to briefly list which facts are relevant and justify the survival of their proofs in (Dembo Zeitouni 2010) with upper LDP.
Definition B.1. A rate function I is a lower semicontinuous mapping X + [0, oc] on metric space X, i.e. ([0, a]) is closed for all finite a. A good rate function is rate function 1 such that ([,, al) compact for all finite a. The effective domain Dr if is I ([0, oo)).
Definition B.2. family of Borel probability measures {u²} } on (X, B) satisfies upper large devia- tion principle (upper LDP) with rate function for all measurable subsets A of N,
supe é log 11 (A) S - -nnf (xx) c-+
(13)
ZEA
A family of Borel probability measures μe on (X, B) satisfies weak upper large deviation principle (weak upper L.P)) with rate function I if (3) holds for all compact subsets A.
For background, recall that (x²) } is said satisfy the large deviation principle with rate function T if in addition to 13) it also satisfies the lower bound
lnf e log μ (A) 2- inf I(z). e-U
(14)
DEA*
Theorem B.3, Suppose A is a base for the topology of X. Then a family of probability measures (ee) } on X' satisfies weak upper LDP with rate function
I(x) = sup AEA,EEA
(- sup lo (A)).
c00
Theorem B.4. (Contraction Principle) F - y is continuous map between Hausdorff" spaces and 1: X + [0, oo] is good rate function. Then
(a) The function 1'(1)= = inf 1 is a good rate function on y; F-'(()))
(b) If a family of probability measures {ur} } on X satisfies upper LDP with rate function 1, then the pushforward measures {nr o F F-1) satisfies upper LDP with rate function I' on y
Remark about the proofs. Theorem B.3 and Theorem B B.4 the upper bound directions of (Dembo)
& Zeitouni 2010, Theorem 4.1.11 & 4.2.1). Their proofs are identical those therein. For Theorem
B.3 note that this direction only uses the equality 14.1, but not (4.1.12) and (4.1.13), in (Dembo
& Zeitouni 2010).
Definition B.5. For families {yr-m) Y and {vS} } of probability measures on metric space Y, where m E N and > 0, we say (ueim) are exponentially good approximations of {v*} if there exist probability spaces (S2, B pr.m) and two families of random variables yf-mm yf with joint distribu- tion Pe and marginal distributions vem ve such that for all 8 > O, the event dist(y "y" > ô is Be 33-maasurable and
lim lip e log Pc (dist(y""",,u) ./) > 6) -00
(15)
e+0
If in addition vmm = T is independent of m, we say {8e} and [v*} are exponentially equivalent.
Theorem B.6. For families {ye.m) and {ve} of probability measures on metric space Y, where m E N and fyrmy satisfies upper LDP with rate function Im then
and {yf- are exponentially good approximations of {v²). 3. If for each m,
(a) {i7) satisfies weak upper LDP with rate function
1(y) == sup sim inf Im
820 m= sEBs(3)
13

---

Published as a conference paper at ICLR 2024
(b) If in addition is good rate function and for every closed subset A S ) inf
T(Y) S sup inf Tm (Y). YEA
m=000 YEA
Remark about the proof. The proof is the same as that of (Dembo & Zeitouni] 2010 Theorem 4.2.16). For part (a), by Theorem 3.3 applied to the topological base consisting of all metric balls Bs(y) V. {r/m satisfies weak upper LDP with rate function
v² I (v) sup (- lim log (Bx(t))) 8>0
c-0
So it suffices to prove T(y) & (v). This done in the proof of (Dembo & Zeitouni 2010 Theorem 4.2.16, part (a)) via the inequality
lim sumsuploog elog (Ba(y)) S
liminf (- inf Im(2)) lim sup inf
(2).
c-+0
m m--000 =EBus(y)
m- =EBas(y)
Hence
I't (v) sup ( (-lim lim sup log (Bs(y))) (lg 8>0
sup lim sup inf [m(2) I(y). 8>0 m-toc =eBaa(y)
c-+
The proof of part (b) is verbatim in (Dembo & Zeitouni 2010
Theorem B.7. Suppose a family of probability measures {uE) satisfies upper LDP with good rate function on a Hausdorff topological space X. And suppose a sequence of continuous maps {Fmm) from X to another Hausdorff topological space y approximate measurable maps F in the sense that for all a < Cc.
lim sup sup (ist(F" F(x)) 0. m-000 (x:1(0){a)
Finally, assume the families ( are exponentially good approximations of another family of probability distributions } on y. Then {r/} satisfies upper LDP with good rate functions I' I'yy) = infp-)(())) I.
Remark about the proof. The theorem is the upper bound part of of (Dembo & Zeitouni 2010 The- orem 4.2.23). The proof stay the same with (Dembo Zeitouni 2010 Theorems 4.2. & 4.2 replaced by their respective upper bounds direction, namely Theorem B.4 and Theorem B.6
B.2 UPPER LDP FOR DEGENERATE DIFFUSION WITH EXTRA CONTROL VARIABLE
In this part, we prove a variation of the upper bound direction in the large deviation principle proved in (Dembo & Zeitouni 2010 Theorem 5.6.7). The notations in this section self-contained and independent of those from other parts of this paper.
Throughout this section will consider the following setting:
U C R is an open domain,
C I be a compact set,
- Rm. U
L(R".R") and h RXXXD R' are Lipschitz continuous
functions,
b and are bounded, and / is bounded on [o] x U x D for some > 0. We will fix e0 and common upper bound I/ for b, a and h respectively these domains.
Consider the following families of stochastic differential equations on (Y,Z) Z) U x R and > 0:
dY b(,, 7, dt + (Y, ddB dz; hi(e, Y di.
(16)
(17)
The main difference of our setting from that in (Dembo & Zeitouni 2010) is the existence of the additional control variable 2;, whose evolution depends on e in a less preseribed way.
We will assume throughout this section, in addition Assumption
B.8 that:
14

---

Published as a conference paper at ICLR 2024
Assumption B.8. For all 0 and initial values (yo, 20) E U X D, the solution (Yf, Zi) to 16 and (17) starting at (yo, 20) almost surely remains in x D for all 0. Definition B.9. Given the functions b, o, the upper bound H on [k], and T The associated path space St is defined as the family of triples (f.g.0 u) with f E C (CO,T) E C ([O,T), and u E w..,2 T), R.) that
ft yo + ! b( [ o(fs...)du.
(18)
Here CLiPP ((0. D) D) is the subspace in ([o, D) of functions with Lipschitz constant bounded by H, i.c. that satisfy
g(s) g(t)
sup ag82
S H.
(19)
And W is the square integrable Sobolev space of first order differentiability.
Lemma B.10. CH C ((0, T), D) compact C ([0.T), D).
Proof. As C* ([0.TT, D) is metric space, it suffices to show any sequence ((k) in CH ((O,T, has convergent subsequence with limit in CHiP ((0, T),
Because D is bounded domain, the g (R)'s are uniformly bounded. As they are also uniformly Lipschitz with Lipschitz constant bounded by H, by Arzela-Asees Ascoli Theorem we may assume g converges in C to some E ([O, D). Then would be H H-Lipschitz continuous as well.
From now on, C ((0, D) will be equipped with the Co topology without further notice.
Definition B.11. Given a function f ([O, U), the corresponding energy functional
Φr := inf (g.a) such that (f.g,)EErr
2 _ L 2d..
By convention, (T( = if ST is empty.
Theorem B.12. Given closed subset A of the metric space A C (OO, TJ, and initial value
yo. for the solution (Y, ZE) to (16) and (17) with initial value (yo, 20), the following inequality holds:
limsupe sup log P(xg (0) (Y, EA inf 2r(s).
c~0
SEA f0=90
In order to prove Theorem B.12 we need more notations to study the Y; while keeping the path Z fixed. For this purpose, we define some distributions of (7i, /e (B) and Y; respectively.
Definition B.13. Denote by 20)T the joint distribution of (Zi, VeB;) on C [0,TT,, x
((O T), ) and Viy. for the distribution of Y in C T). U), where (Y;, 77) is solu-
tion to 16). 17 with initial value (y0,20) 20) at t O. Write X for the distribution of VeBi
(10, TR, Rd
Definition B.14. Denote by MG the space of all probability distributions μ C (00, 27, D) x
([0 [[, (Rd) whose projection to the second coordinate is X.
For all initial values (yo, 20) EUXDYY, Yr. Zi are progressively measurable processes. In conse- quence,
(. E M,
(20)
Lemma B.15. The function
Lel 2d iuue 2 18)
[(g, u) =
2
20
otherwise
15

---

Published as a conference paper at ICLR 2024
is good rate function CH (I0, D) x ([0, Moreover, any famity of probability measures {ur} on C (10, 1|, D) x CO ( 1]. IR)) such that μº E M satisfis upper LDP with rate function I.
Proof. We first check that I is a rate function. That is, it is lower semicountinuous on CH (I0, D) x ( R.). It suffices to show that that I(g, u) S lim inf I[g u (k)) if
k=4000
(g(h), () (g, u) C norm.
By Lemma 1. g is in C ([o, D), and thus by definition [(g. u) = =I((u) in this case.
It was known by Schilder's Theorem (Dembo & Zeitouni 2010 Theorem 5.2.3)) that the function
{ !
ur /
if (1,,
To(u)
oc
otherwise.
is a good rate function, and in particular lower semicountinuous. Thus
lim inf (gghh, (k))) lim inf lo(u(*)) 2 Io(u) [gg, u).
Thus is a good rate function and we want to show that it is good, i.e. T-1 a)) is compact for all
finite a. Remark that ((0, al) = CH ([0. 1], D) x 1o (00 [[0,a]) and the second factor is compact To is good rate function. Note that (u) T(g u). So suffices to know that CH 1), D)
compact space in C topology, which is the assertion of Lemma BB.0
Finally we need to show that 13) holds for {pr} and I. The same inequality holds for {X} and Ios
ie. for all Ao S ( 1], ).
lim sup e log (AAo) - inf Io(u). e-0
uE
For all measurable set A C CHP ((0, D). x (=0 (], R),),, let Au denote its projection to ((0, 1]. Rd). Then
lim sup uf (A) S limsupp log n (CH ((0, 1]. D) x Ao) -0
e+0
lim sup e log (AAo)
c+0
S - inf Lo(u) S inf [(g, u). (g.)ed
mE:Ao
In the last inequality, we used the fact that Ao contains the projection of A and that [(g, u) lo().. This completes the proof.
For all (g, u) E CLip (1O, D) x ((0, ))), denote by Y (g, the solution on [o, 1]. with initial
value y ((g.)).0 = to the stochastic differential equation
dye u)) = (.VG..)..: di + (V (g.a),s as 9t )
(21)
In addition, given an integer E N, we also define (g.a). the solution on [0, 1]. also with initial value yo. to the following stochastic differential equation
dy b( (a.w),t (g.4), [mc] [ :9t )d + a (g.x), [me] - dd..
(22)
We emphasize that Ye (g.n), and (q.4), are deterministic once the pair (g, u) are given and all ran- domness comes from u.. "
Lemma B.16. For all any 8 > 0 and all initial values yo E D,
lim simsupe sup log (g.)(u))))u((( sup (. Y Y..,), 5) -o0. c-+0
MEM;
tE(0,1]
16

---

Published as a conference paper at ICLR 2024
Proof of Lemma B.16 Fix 8 > 0. Let A c.
(g.u).t
Y
Y (g,u),, for any p 0, define the
stopping time
T TSmPP : (inf (g.u).
yem (a.u),
2 pt 1).
[md
7e.m-e depends on (g, u). but skip it simplify the notation.
The process A e. (g.a),. satisfies the SDE
6m (0,))r
= bgdd dl + vear duc
with coefficients
b = b(ycm
(g.u). [mm] : b(Y (.
of = o (Y
9t)
o (.).
(g.u), [msi
By the uniform Lipschitz continuity of b and and the boundedness of Jr E D, there is a constant C' such that for all S T6.m..
max 1, 1.
C(I (aww), + p).
By (Dembo & Zeitouni 2010 Lemma 5.6.18), for all E (0, 1),5 0,
p²
sup log P (g,u)aal sup te(0,rr.m..)
A
(o 2 0) < K + log p2 + 82
a.m
HEM
where K is constant independent of 8, e, P. l and m. Then
lim sup lim
log P (a.))ma HEMi
sup
sup IE refom.
A (ga)ael 2 8) = -CC
p-0 m21 c-0
Since the event (sup
/ t.m (g.u),
2 5} is contained in the union
< 1)} U
sup tE(0.rs.m.p))
(ou)) 8}.
the lemma is proved we show for all p > 0,
lim lim sup log (a. sup |ye.m yem (g.u), Imil 1 p) -00
(23)
(g.a).t
m-000 e-0 MEM;
sE(0.1]
To prove this, recall that |b] and |o| are bounded by constant H, and
(a.w).t
pem (g.a),
imme
S
+ VV
max
sup
1u à.
1
m
osas
Therefore, for > H
sup log Plg.avpp(
sup
(g.u),t
rom (g.0s.
Jimel 2 p)
NEM
0stš1
p
SmIP( sup
|us|
CCC
US 4dme m(p- /2deC
This guarantees 23 and proves the lemma.
Proof of Theorem B.12 First of all, notice that one assume without loss of generality that 7 by rescaling the time interval [0, T] [0, 1]. To see this, notice that the rescaled Brownian motion
Bd, is equivalent to vTB4 on t E [o, 1] and we make the change of variable ur: = VTv. accordingly. then
2 ju dt
1 1 wrel 2adtt [''
L.
11 dt.
2 dr
T
We will only deal with the [0, 1] interval below.
17

---

Published as a conference paper at ICLR 2024
Define maps FT, F: C (I0, D) ([O - ([0, U) follows.
For FT,, the image fm = Fm (g. u) satisfies fo" = yo and
{ 6(f -o) ds f of g. g.)
St
(24)
on t [ [+ }, 0.... 1. m
For F, the image I = F(g, u) instead satisfies fo = yo and
fo b(sss L'/ o
(25)
for [0, 1].
not hard to check by Lipschitz boundedness of b, a, and the compactnes of that F and send CH ([0, D) C ([0. 1]. RR)) to C (Io, 1], U). Moreover, (21 and (2 can be reformulated
as
Fm (g. u) and Y (g.w),t F(g, u) respectively. yem (g.)),
(26)
For u), (g', u') E C ((0, 11, D) x ( 1], (Rdd) let
- = Fm (g. u) - Fm (g': u')l.
TT is a function of t, and by the Lipschitz bounds on b and o.
OT < C(eT h ) + ll(., u) (g', lo)o) m
sup KE mo ++1
Since (0) 0, we derive the continuity of the operator F by iterating the argument with 0, I,.., m A similar argument guarantees the continuity of the operator F.
By 20 and Lemma B.16 for all family of initial values {26 E D}e>a.
lim lim e log P (g.n)) (6v c-+0
(sup rE(o.1]
|Ymm Y (a. 2 8) -00. -00 (anw),t
In other words, the families {Hr (1o,))))1 (Fm)-1} are exponentially good approximations of the distribution of V (g,w).t with (g. u) M (101))))11
Thus if we can prove: for a < 00.
lim
sup (0 Hl desa
|F (g, u) F(g, w)llo O,
(27)
m m0000
then by Lemma B.7 and Lemma B.15 the distribution of (gg.),. with (g. u) M(yo.st).1 satisfies upper LDP with good rate function y + inf P(g.)))= [gg u). which is exactly (v). Because for each different value of c, 25 is arbitrarily chosen in 1 and the distribution of V (g,u),, for (g, u) ~ M (1 coincides with the distribution "(m.e6).1 of Y; with initial condition (Y,, 2i) = (Vo, =f)
in (16). 177. this exactly yields the statement of the theorem.
remains to show 27 Let
m
p u) - F(g, u)l.
Note that if
[ 1 Pd S
holds, then for pm pm u) and f F(g, u) given in
24 25
sup max (ITT te
[ma] life f [mtl
< Co
m v
Coa
1
(28)
m
18

---

Published as a conference paper at ICLR 2024
for constant Co by the bound the coefficients and Sauchy-Schwaars inequality. Write yo for the right hand side in (8). Then by Lipschitz continuity and 0, for some other constant C,
1 1 fil
1/" (olFT (mflmml9e)) b(fer (.)) S/" (ol Imsl (mmm)ge)) offo g.) údds
Offmai - -fa + +f(
(((00) c1 [mal - fal /d,
C(A + m1 + (ú.ldds
By Cauchy-Schwarz inequaall for all t E [0, 1]
(A
(A m / (1 + l
<C² / 2((A,)) + (1/m1) )) / 2(1 + | 2d
<40 (1+ / (AIr)
(1mm)2).
Thus by Gronwall's inequality, (Ar")² < 40 + a)e (Lt)) The equality 27 follows by letting m + oo, This completes the proof.
The following a strengthen version of Theorem B.12
Theorem B.17. Given closed subset A of the metric space A C ("O, 11, and initial value (ye, 2.) x for solutions (Y Zz) to (16) and (17). the following inequality holds:
lim sup sup clog P (1m (Y E A) inf
Dr(f).
c-0 meD
SEA
fo=y.
Proof of Theorem B.17
As in the proof of Theorem B.12 we can assume T = 1.
suffices to prove that for any family of points 36 - y. as c - and arbi- trary z D. the family of distribution {V( is exponentially equivalent (((g..))))) in Definition Write
By Theorem B.6
Y (( Y
(5)+
A =
2 (1b - D (v.,=)),
Then A satisfies
dA a,dt + d,
where
b(Y: ((
Z
(6).)))).) ) h(Ve ()-1z)))e: Z: (onn)))))
b(Ye (1 Zi
of hiYY
2e (066)))),)
(161=6))e:
and
B8 = a (Y(() ()) Zi
(Ye (V of)eer ((y.es))d) )
Moreover,
A6 = y6 - 3/ - 0 as +
Remark that the coefficients af, 3f are progressively measurable processes with respect to the filter generated by the Brownian motion {B.). Moreover, Lipschitz continuity of b, h, for some constant Co.
max(lau), (3E)) < Co/Af1-
19

---

Published as a conference paper at ICLR 2024
By applying (Dembo & Zeitouni 2010
Lemma 5.6.18) with 71 1, we know that there is another
constant C' such that for all p 0, >0, 0,
p² + 16 ya 1
e log P( sup [0
A 2 6) < C
log
0² + 82
By letting p - first and then 0, it follows that
lim sups Pog P( sup
| 2 8) -00.
2E[0,1]
This shows (v(6.==),1)) is exponentially equivalent (V)...)), J. which suffices conclude the proof.
Corollary B.18. Given closed subset A of the metric space A c (10, U) and compact set K C U. for solutions (Y;; 7Z)) to 16 and 177. the following inequality holds:
lim e log e-+
sup EKXD
Por: (0n) E EA
inf Pr(f).
fEA fock
Proof. Let M E [0, 0o) be finite value strictly less than inf SEA ( [0, oo]. By Theorem loEk
B.17 for all y K, there is value ey such that for all yo Be, (y) and 0 < <<<e Ey.
sup e log P(xg (;)=(vo. (Y, E A) < MM
#0ED
Cover by finitely many balls of the form BeG (wi). then for all < & < min; Eys
sup sup e elog P(rg (1)= (YE E A) < - MM. k =0E
The proof is completed by letting M - inf SEA (f))..
foek
B.3 EXITING TIME FROM BASIN
In this section, it will be assumed, in addition to Assumption B.8 that:
Assumption B.19. There are C2 function L U + [0, oo) and a bounded open set V C U such that:
(1) (b( s for all (v,2) V x D and the equality holds if and only if L(y) 0:
(2) L is strictly positive on OV.
Write Vi L- L-''((,, 9)) n V and choose qu sufficiently small such that Vg lies in the interior of V. In particular, for all g E [0, go). Llav, = q-
Definition B.20. Suppose 0 < q < Q S qu- For a solution (Y,, Zi) of I6 17
with initial value
in VQ x D, denote by T the first time Ył hits V, U UVQ. and by To the first time Yi hits VVq.
Our goal to prove the following main theorem:
Theorem B.21. Under Assumptions such that for all,
B.8 and B.19 for all q Q < go there exists IQ >
JQ =
lim Oo
sup
P Ply zg) (00 20) (T)
The following lemma is analogue to (Dembo & Zeitouni 2010 Lemma 5.7.22)
Lemma B.22. For all 0 < < Q' < <Q S qo, the stopping time Th satisfies
lim c+0
sup (VoxxD)
For (xx622))( ( < oc and Y : E V,) = 1.
20

---

Published as a conference paper at ICLR 2024
Proof. Consider the solution Y the deterministic flow
dY," b(Y, =o)d/
(29)
starting at yo- By Assumption B.19 L(Y?) is decreasing and must converge to 0 as t - 0, and Y,
remains Vo. Denote by T the first moment L(Y?)) reaches z Set
8 =| min 9 Q 2,0-% Q'
0.
2
Notice ( dist dist z{), , 8 implies
1
Q L(y0)
L(Y,) L(x")) S min
sup
4
<E[U,2]
Since I(VY) & I(Y)) = Q' and I(YP)) = 2. this guarantees that
L(Y - L(Yo) + 1|
and
I(VH)= S !
+|[T[lcar <q
would in turn follow that 7. < f < 00 and
Y 7.0
E Vg_ Hence it suffices to prove
lim P (Y;:2))=(M010)) c-0
(Y 70, (YP 20)l 26) 1. *E(O.T]
sup
(30)
For simplicity, write J = (Y 2,), (YP, 20) Let M < oo be a common upper bound on the C and the Lipschitz norms of the maps b, and h over the domains E [0, 1]. y E Vgo 2 E D. Then JÁ 0 and
2MJ + vC L'' o(Y Zy)
J j0
By Gronwall's inequality,
sup J Vee sup rE|o,7]
J. aVa Z 23)d
rEJo,T]
Therefore, in order to make suprejo,7
Ji 2 8, we must have
/ o , 2 2 de te- -2MP
sup (E|0.2]
and thus
f OY )
de - 2M1T
sup [E[o,T]
>
must hold for at least one of the row vectors of a.
The process G = So (Y k3)d is continuous martingale with increasing process (GE) S M almost surely.
Recall (G),, is defined l of(Y,, ds) By the Burkholder-Davis--uuddd inequality (see e.g.
DDembo & Zeitouni 2010 Chapter E)),
B(yg.2))= 21)==(0.z0)
sup
IG I C. (2)= (G CM
(E(o.f]
for absolute constant C. And thus by the Chebyshev's theorem
L/' (Y, 23)BB
P F(Xg -2)==(x..z)
sup tE(0.T)
2 J - 22M
SCMT( _e - --- 2MT -2M
21

---

Published as a conference paper at ICLR 2024
From the earlier discussion, after summing over all k's, obtain that
sup J 2 6) S CMPTG je - :<(0.7)
P( (2)))))))..0
(31)
We deduce (30) from 31 by letting + 0, which concludes the proof.
The following lemma is the analogue of (Dembo & Zeitouni 2010, Lemma 5.7.23). Lemma B.23. For all a > 0 and bounded set K U, there exists constant To = (o a, K) > O such that
sup e log c-0
sup KXDD
P( (x)022))=(10 ( sup tEJo.Tol
(Y;
Yo. Z =0)1 2 8) < -a.
Proof. Without loss of generality, assume E, d E [0, 1] with å fixed and e varying. Let the stopping time 5 be the first time such that (Y; yo. Z - su] 2 5. Then for every t [0.
16 2,)
max Bg(K)xD
[6] + 1
( - S max
o + ||o||
Bx(K)xD D
[n(s, ZE)
max [o.))xB((K)xD
Ih| + | Lip
and thus they all are uniformly bounded by constant M. For all S t S min((" 4M ). !Zj - =ol S
i and YY, vol s I (Y Zj)dBBB Hence for any To I aM
we have
P (V) 15 2 To)
2 6 ) L a(Y: 2.)
(c 2 To and Ve sup tE(D,T]
<Par (X: (x)
As in the proof of (Dembo & Zeitouni 2010, Lemma 5.7.23), it suffices to consider each row vector o of ar The stochastic process J ax (Y; Zj)dBB is equivalent to B,g, by the time change theorem (see (Dembo & Zeitouni 2010 Chapter E.2)) where Thé = St o Zz)ds. The function T is increasing in and almost surely Thü S M2 To 5 as the C'9 and the Lipschitz norms of b, a, Lemma 5.2.1). 1),
are bounded by M before the stopping time Çτ. By (Dembo & Zeitouni 2010
(5 To and Ve sup tE(o,Ta
J ox(YY, Z.JdB!
5
P(y (6))2)))(mo.00))
Z 2d
SP (x) c 2 To and Ve sup bE[O.Ta]
2 TE 2d
B² 2 8
<P (y6.2))=(vvoo))
2 To and Ve sup re(o,
2d
<(Yy (16)2))-(N0-1)) (1o. (Ve sup 72(0)
B₁ 2 2d
/
Summing over different i's, we conclude from the two inequalities above that for all To [0, 1
Poy (3)= (S 2 To) < d 4e (2) (22MM))) = 4de 1 2
It now suffices to take To min( TM 5 M.
8²
Definition B.24. For yo. E U and T > O, define an energy cost Vr v1) by
Ur vn) innf { { T
2 j. ids: (f.g,)) E Sr, fo = yo. fr = y1
innf ( ((0,,)7)
22

---

Published as a conference paper at ICLR 2024
The following lemma claims there is minimum energy cost required to increase the loss function between different levels.
Lemma B.25. For all Q 4o
inf T20 (vo Ve.. xoVo
inf
Vr (30:0.911 >
(32)
Moreover, there is constant T. >
such that inf
inf
VT v1)
(33)
T T24.0 y0.3n EVo.
Proof. Suppose for yo, 1/1 E Vi ov and (f.g,u) Sr with fo =yo- fr y1- Then by
aL(S)) =v L(fidfft
= -(V ((fa)) Is (V
I(y)) = Von..) (ty) < for all y E Vo and = E D. particular, this also shows VL(y) + for all y Vo Since both Vo\VV and D are compact, There exists positive constant x(q, Q) > 0 and n/g, Q) 0, such that (V L(y))b(y, 2) xx))( and [V(()) for (y. E (Vo V1)) D. Integrating from 0 to T, we get
Because q > by Assumption B.19
L(yi) I(yo)
-I(fr) I(Fo)= = (T
dI((f)) {
= 1." ((VT ((TT))))6(Gon.o.ll + ((T)) )dus)
S n: {'' (VL(f)) 2ds
(34)
(VL(s)) 3 (o (s
5 !'" (VL(s)
S (VL(s))
|| ( (% | (da))ds)
In order to show 32). note that if yo E Vg and yn E TVg. then L(y1) L(vo) = Q > 0. By 34 and Cauchy- Schwarz inequality,
Q-95 - K ''" dL(.))
J (L(())) "alllon (f,² linPas)
(35)
(" li./'ds.
4k
Since is continuous, || ||ollc>
co. It follow that
/
T
1 dss > 2x(Q 9)1اا >0
This provides positive lower bound for v (yo. v1) that uniform for > 0, yn E Vg and yn VVg. This proves 322
23

---

Published as a conference paper at ICLR 2024
We now prove 33
Note that for all yo, 311 Ve, L(n)) L(v0) 2 2 Q. By 34)
So |ua (ds)
(R() (=(( [L((a))dss) (Q }
-000-) [VL()a))as) $)
|||0||| [| vQx
2 2| x (KTTI - (Q - qim "T
The last expression is uniformly positive for T' 2 Tq :- 10Q--gx--1p Thus T v1) is uniformly positive over the region given by T 2 q, a yo. V1 VV.. This proves 33
Lemma B.26. For all 0 < Q S qo and initial value (yo, 20) E VQ x D, the stopping time q.
satisfies lim lim s log c-1
sup (oD
Por: (0) (T > T) = -00.
Proof. Note that T. T implies Yf E Volv S Vo for all []; or in other words, {Y, hejo.TT E (I, Vg)
Therefore, by Corollary B.18
lim lim
sup (V
(0) T)
T T-x c-0
S inf
o(f);
[EC"([0.T).Va VoN)
and in consequence it suffices to show
((f) = oc. T-000 (o.7
lim
inf
(36)
Assume, for the sake of contradiction, that (36) is false, then for some M < and any all k 5 N.
there exists f& E C ([0, q), VV)
with (Jkk) & M, where Tgo0 is given by Lemma B.25
After breaking J& into l segments subintervals of length Tq it follows that there exists J CO Tq, Voy VV) such that ((f2) <
By taking limit in C norm (which permitted by Lemma B.10) and using the lower semicontinuity
of the good rate function 4, there exists f (² T. T. VQ' Vgy) with (f**) 0. This implies
(y y1) and thus contradicts the inequality (33) of Lemma B.25 This completes the proof.
Denote by I, > O the left hand side in 32
Lemma B.27. The solutions to 16 and
satisfy
lim lim sup log 40+
sup (00.20)EV2. EV,
Porg (vo.z))
(Y:. 7. E eVV2) < lim Iy.Q. 2-00
(37)
Proof. Fix an arbitrarily small 8, and let Iq == min(limg-o) Iz.Q 8, 1). Note that the right hand side in (37) always exists because Ig is a decreasing function in q by construction. In particular,
Q 2 Iq.s when is sufficiently small depending on Q and d.
By Lemma B.26 there exists large T. T. T.(q, Q) oc such that
lim sup c log 00
sup Eo),
Por, (v0. (T > T.) S -ĪQ.s-
(38)
In addition,
inf (o. foEV2 suPreja.r.] 1S)2QQ
((f)) 2 Iz9 2 IQ.51
24

---

Published as a conference paper at ICLR 2024
and thus by Corollary B.18
lim sup c log
sup (00=20)EV2,>
P(y: (1)0.0))) sup L(Y)) 2 Q) S -IQ.6- t0
(39)
c
Note that the event {
E 0V2) is contained in the union of the events {T > T.} and
{supre[u,r.) L(X;) 2 Q}. Therefore, we obtain by combining 38 and (39) that the inequality
lim s e log c-0
sup E2g D
Porg ()) (Y 1. E @Vo2)= S IQ.
(40)
holds for all sufficiently small q. The lemma then follows by letting å + 0.
We are now ready to establish Theorem
B.21
Proof of Theorem B.27
Choose
Q to be sufficiently small. Define sequence of stopping times recursively by letting
80 TO
< 01
< T
0u : 0,
Tm == inf
Om Y E V (VV]),
On = inl Tm-1 Yi E 0V
By Lemma B.22 all these stopping times are finite almost surely. Write Ym Y which is Markov chamn
0 for all 0 q Q and is decreasing function in g. Let Iq = 1/ lo I, 0 and fix 0 < LīQ. By Lemma B.27 if we fix sufficiently small q, then
Recall that by Lemma B.25
Ig
simsupclos log <-0
sup (o) x
Poy (17.z))=(a (v)= (Y (VQ)
Iq
3a.
By plugging in (Yg,, Z5 as the the value for (y, 2), we deduce that there exists > 0 such that for all 0 < 6p and m 1,
sup (2o.2))EVQxD
llgg (r (0) -mm) 5 S--S-Q 20x
(41)
On the other hand, assuming 6o is sufficiently small, applying Lemma B.23 with = Iq + 2o,
8 = 1g and K VQ yields that, for some fixed To > 0 depending and Q and all (0, col.
sup (V
elog Ip Pog (0)0
Tm-1 S To)
1 y, 2 - 2)] 2 (
sup (g.=)EVoxD
elog P (yg (8) sup |(Y tE(o,To)
(42)
< - Iq= 2a.
For given M, the event (T) & MTp} is contained the union of the events JM (Té = Tm)
and U (0m - Tm-1 < To Combining this fact and the incqualities 41 44
yield
sup
Fry,,2))(0n S MTp)
(VV
sup (V,
(1) To) Σ
sup EV,
P (20. = Tm)
(43)
M
+
sup (o,,>))
Piy P(yg (0))))))=0000) Tm-1 E TT)
Fg -
sup
P P1yg:23)=(p0.00 (76 T0) 2MM
(V,
25

---

Published as a conference paper at ICLR 2024
Set M = 17o 30+0
Zg
which makes MTg 2 6 < for all sufficiently small 6. Moreover as e + O,
by the (p (o) VV, VVVxDD XD P (g)== (T = TD) + 0 by Lemma B.22 and 2Me. 1o
choice of M. Thus
2q P(x; (2)=220)))0000 (To Ser)
sup (o V
sup
(0) S MTo)
(g) (o, X
XD
which tends to 0 0. This completes the proof of Theorem B.21
B.4 PROOF OF PROPOSITION 4.2
Recall that the movement of 2 stays in T i.e. |xx]| [R_., R4] and can be characterized as the new model (6)
After applying time change
T= (m1))b
and writing
2t = n7 = An |x,1 !
get the following system of equations (X,, )rr g0-1 x R+:
dx V + (n $ 7/ dBB
(44)
dyr (n (-45r + 2Tr E(XT))dT
(45)
To deduce 44 we used the standard fact that for any a > 0. a B B and B are equivalent as Wiener processes.
Note "r stays the fixed interval [7-,5+] = [Arp- RR Xuj [R as long as E
As there are only finitely many basins, we may fix an index i without impact Proposition B.28 In addition, without loss of generality let us assume
on I For our purpose, it might be better to measure the distance to I'd on #.po where Pi is sufficiently small and in particular
#opo are disjoint for distinct j's. For 2 O, we will write V = {o E Up L(x) < q} and V#.4 (x U1 #.po L(x) < 4)
By Assumption
1.2 (ii), one may manipulate Po. 90. PL, Q1 (Pu > Pl, 40 > 91) such that
U S VV.
V S U
and thus reformulate Proposition 4.2 as
Proposition B.28, There exists > such that if go > q1 > are fixed, but qı is sufficiently small
compared go. then in the regime v O()) < O(1),
X remains in Vi for all i E [0, ((n))))) $1) 1.
lim n 20) x19-17+
sup
P P(Xo.1))((00=0)
The convergence is uniform with respect to (yo. 20).
Proof. In order to apply Theorem B.21 with domain V and control domain D = 15-, 74], we first make the following remark.
Theorem B.21 is in the setting where the domain is an open neighborhood in R", while our current Vi is neighborhood in the sphere. This is not problem because unless L is constant function, ri is proper subset of the sphere $d-1 And V is also a proper subset in gd-1 for small One can then change coordinates and identify Vi with subset of the Euclidean space.
This converts the problem to the equations
with the following dictionary: c (n))) X and रr play the roles Y,, and Z respectively: b(y,2) = -2 Vv)) oy. = 2 =to(y); and h(e, y, =) e(-44 42 + 2 Tr E(y)).
26

---

Published as a conference paper at ICLR 2024
It remains to check Assumptions B.8 and B.19 We start with the latter. The strictly positivity of directly follows from the construction of and the neighborhood V The property (1) in Assumption B. holds because b is negatively proportional the gradient V of L with respect to the spherical coordinates.
Unfortunately Assumption B.8 doesn't it automatically hold trajectories may escape from Vi.gu In order to adapt to this case, we smoothly modify the value of a on Vi, so that it remains unchanged
on Vi.q and vanishes near 0vi 1.go" Then near the boundary, the SDEs (16).
7 become determin-
istic. Because of Assumption B.19 (1), the trajectories do not escape from V 1.000 Moreover, we
know that T remains in D = 7-1 This verifies Assumption B.8 for the modified model.
We conclude by applying Theorem [..1 that, for some fixed I > O, for all initial positions in V1.9 x [+], the probability that a trajectory (XT,5) 7r) (with respect to the modified model)
leaves Vi.go [5-.7+] before = e कx cO((nA)=2) converges as 1/> 0. The convergence is in addition uniform with respect to the initial position.
Since such modifications only take place outside Ví,,, 7+], the same statement also holds for the original model. As Vi.. C Vo. we obtain the statement Proposition B.28 after reparama-
trization of variables.
C UNIQUE KATZENBERGER LIMIT INSIDE EACH BASIN
guarantee that, after discarding an exponentially small subset of random incidences, the trajectories of (2) stays inside the basin that contains the initial position for exponentially long time ((nA) We now justify Proposition
The results from $8.2 stated in the form of Proposition 4.2
4.3
We restart
② from an initial point, still written as To by abuse of notation, in some U We #.p.
now apply uniform approximation theorem, which a stronger version of (Li et al. 2022a Theo- rem 4.6). We can prove this uniform approximation result because we can strengthen Katzenberger
1991 Theorem 6.3) to be compactness theorem, uniformly with respect to the initial points of the SDE. To describe this result, let (S2 F {Fr 1/20, P) be filtered probability space, Zn an R and A, a real-valued cadlag aadaptdd nondecreasing process with A, (0) O. Let on U M(d, e) be continuous with on uni- formly on compact subsets of U. Let X,, be an Rd-uul cadlag {FF") -smimarrtingale satisfying, for all compact K U.
valued cadlag (J() -semimartingale with Zr (0)
(t(() = (0) + ( o(Xn)dZn + S." -VL(Xx)dA
(46)
for all & S (n( where \n( inf(l 2 0|Xm((-) $) or Xa(1) K is the stopping time of Xn leaving K.
Theorem C.1. Suppose (0) E U, Assumption 3.1, 3.2 and Condition B.2, B.3, B.4, B.5 from LL et al. (202aa) hold. For any compact K U define p.(K) inf (t $ K or Y,,()) &
K}, then the sequence ((YHil()) (H,,(K) , ,Po(K) (K)) is relatively compact in Daaxe (0, oc) 10, oo).
If (Y, Z. p) is limit point of this sequence under the Skorohod metric, then (Y, Z) is continuous semimartingale, (()) E T for every! t 2 Das., p inf 2 (t) & K} a.s. and (t) admits
Y(t)= -Y))+ { MAN
(Y(Y()))
d
(47)
+ 2 1,7 k,1=1
o, ( (s).
We will present Assumption B.3, B.4 and Condition B.5, B.6, B.7 and B.8 from Li et al, 2022a in Appendix B.
The main difference of Theorem to (Katzenberger 1991 Theorem 6.3) is that we allow the initial point X,, (0) to vary within U.
27

---

Published as a conference paper at ICLR 2024
Theorem C.2. Let the manifold I and its open neighborhood U satisfy Assumption 3.1 and 3.2. Let K C be any compact set and In,0 E K be a sequence of initial points. Consider the SGD formulated in 44) where Xmm (0) In Define
Yo (t) X no (t) (X m. (0), Ar (t)) + (X,,,,
and 1o, (K) min(t E NIY, (t) & K]. Then the sequence (YH,.(K) Zz Pp [K)))n21 is
relatively compact in Dgaxe [0, x)) x (0) If (Y, Z, U) is a limit point of this sequence, it holds that (t) I' a.s. for all 2 0. μ 2 inf(t 2 0 (t) $ K} and Y (t) admits
(t) = //-0
AAN Σ
(X(X))))
(,001((6)66(XX))))))X))) (oyds (48)
+==1
where (W(s) is the standard Brownian motion.
Proof. The proof of Theorem follows how (Li et al. 2022al Theorem B.8) was proved by using (Li et al. 2022a| Lemma B.6) and the standard Katzenberger's theorem (Katzenberger, 1991] Theorem 6.3). One difference is that here not all trajectories stays inside one basin. However, we
claim that the probability that trajectories escape the basin goes to zero when 7 tends to zero. Once this claim is proved, Theorem C.2 is an immediate consequence of (Li et al. 2022a Lemma B.6)
and Theorem
To prove the claim, adopt the same idea as in the proof of Theorem 4.2 in Appendix B.4 Al- though trajectories may escape from level set Vi.go we can smoothly modify the value of on the
closure Vi.a, so that it remains unchanged on Vi and vanishes near the boundary avi. The SDEs become deterministic, and thus the trajectories of the modified model do not escape from the boundary avi... Now, as y7À tends to zero, by Theorem B.21 the probability of a trajectory of the modified model leaving avi.g, x (R-, 7++ before a fixed 7 tends to O. Since such modifications only take place outside 0Vi, x [--, the same statement holds for the original model. This finishes the proof of the claim."
L
Proof of Proposition 4.3 The above uniform version of the Katzenberger's ' theorem guarantees that, starting from different initial points in the same compact neighborhood of the basin, the distribution of trajectories associated with ② is still close to that of the Katzenberger's SDE (47). By Proposition 3.1 the latter is mixing towards a unique equilibrium z/. Note that even though in Chapter 3, we have only proved it for one basin case, Theorem 4. shows that with large probability, the trajectories do not escape from the basin. For those trajectories, they satisfy a modified SDE equation like before, so that all trajectories do not escape from this basin. At this moment, can directly apply Proposition
3.1 It follows that within any polynomial time window under consideration, the distribution of trajectories associated with (2) are also mixing towards v². This proves Proposition 4.3
Assumption C.3. (Li et al. 2022a Assumption 3.1 Assume that the loss L Rd + R is a C function, and that I is (d M) -iddimensio C² -u-ummanifold of Rd for some integer 0 S M < d, where for all E T, x is local minimizer of L and ronk( (V))(())))())))))) = M.
Assumption C.4. [Li et al.) 2022a Assumption 3.2) Assume that U is an open neighborhood of satisfying that gradient flow starting in U converges to some point in 7,, i.e., Vx E U, ((x) P. (Then D is C² on U by (Falconer
1983).
Condition C.5. Li al.
2022a
Lemma B.2) The integrator sequence (An)a2 is asymptotically
continuous: sup (An())
A.
(1-) = 0 where (n((-) = limm-st- An(s) is the left limit of An
tt
2022a fast: VE 0, inf,200 (2000 (A ( +c)
Condition C.6. [Li al.
Lemma B.3) The integrator sequence {An)azi increases infinitely An(t)) + oc.
Condition C.7. (Katzenberger T > 0, as 00, it holds that
1997 Equation 5.1), LLi et al. 2022al Lemma B.4)) For every
sup OCISTAA,
(t) 112 0.
28

---

Published as a conference paper at ICLR 2024
Condition C.8. (KKatenberrge 1991
Lemma B.5)) For each n 2 1, let Yn be a { -semimartingale with sample paths in Dga 0. oc). Assume that for some 5 > Oallowing ö = and every n 2 1 there exist stopping times {-e (m 2 1 and a decomposition of Ja(Yn) into a local martingale Mn plus a finite variation process Fn such that P( S m] 1/m, [[M,](t A xx) + Tenrg (Fa) (Fa) uniformly integrable for every 2 O and m 1, and
Condition 4.2), (Li al. 2022a
lim limsup Pl sup (F.)( > =0 7-0
n-020
OS1ST
for every > 0 and >> 0, where (-)-) denotes total variation on the interval [0, 4].
was shown in (Li et al. 2022a Lemma B.6) that for SGD formulated in 446). the sequences {AAn h221 and [Z |n>1
satisfy Condition C.5 C.6 C.7 and C.8 And the landscape of L satisfies
Assumption C.3 and C.4 Thus the Katzenberger theorem holds in our case.
29


Reviews:
[{"review_id": "Dz8jbIEAgO", "reviewer": "Reviewer_jkii", "summary": "This paper proves the fast equilibrium conjecture for SGD on neural nets with normalization layers in a more general setting than previous works Li et al., 2022c and Wang & Wang, 2022. Specifically, it shows that the conjecture still holds without the unique basin and analyticity assumptions made in Li et al., 2022c. The theoretical results are further supported by experiments.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper is mathematically solid. It extends the conditions for the fast equilibrium conjecture to a more general setting, making a good technical contribution to the community.", "weaknesses": "1. Referring to the second experiment as \"stochastic weight averaging\" (SWA) is inappropriate, as SWA averages the model parameters at different iterations along the same trajectory. Conversely, the approach in this paper averages the parameters at the same iterations from different trajectories.\n\n2. Missing discussion of related works, which \n\n   - argue that the iterates stay in the same basin for a significant amount of time when starting from the same initialization. \n     - Frankle, J., Dziugaite, G. K., Roy, D., & Carbin, M. (2020, November). Linear mode connectivity and the lottery ticket hypothesis. In *International Conference on Machine Learning* (pp. 3259-3269). PMLR.\n     - Gupta, V., Serrano, S. A., & DeCoste, D. (2019, September). Stochastic Weight Averaging in Parallel: Large-Batch Training That Generalizes Well. In *International Conference on Learning Representations*.\n\n   - Also analyze the dynamics of SGD / Local SGD near the manifold of minimizers\n     - Damian, A., Ma, T., & Lee, J. D. (2021). Label noise sgd provably prefers flat global minimizers. *Advances in Neural Information Processing Systems*, *34*, 27449-27461.\n     - Gu, X., Lyu, K., Huang, L., & Arora, S. (2022, September). Why (and When) does Local SGD Generalize Better than SGD?. In *The Eleventh International Conference on Learning Representations*.\n\n3. The paper is abstract in its current form. It would be beneficial if the authors could provide specific examples where the removed assumptions may be too restrictive.\n\nBTW, I did not have time to check the proof. So, my final rating will be influenced by the evaluations of the other reviewers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ofexroW8MT", "reviewer": "Reviewer_yR6a", "summary": "This paper provides a strengthened proof of the fast equilibrium conjecture that was proved in the previous works (Wang &Wang (2022); Li et al. (2022c)) by removing the non-generic assumptions of the unique basin and that the set of minima is an analytic set. In order to remove these additional assumptions, this paper mainly adopts a purely probabilistic method rather than the spectral analysis that was used in the previous works. \n\nToward this goal, this work shows that trajectories would not escape from the initial basin in exponential time.", "questions": "1. Will the noise structure affect the convergence rate?\n2. Would it be possible to achieve a similar result if $L$ had a homogeneity degree  $> 0$?\n3. Are all the remaining three assumptions essential to prove the Fast Equilibrium conjecture? Would it be possible to even weaken the assumptions?\n4. Minor:\nIn the first line in Section 4, \"Assumption 1.3. (ii)\" -> \"Assumption 1.3. (i)\"?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This paper is clearly organized and well-structured so that it is easy for the readers to grasp the main contributions of this work.\n2. I like the fact that this work provides solid and well-supported arguments (Remarks 1.4 and 1.5) that Assumptions 1.2 are natural and Assumptions 1.3 are non-generic. These motivate removing these less natural assumptions well and signifies the contribution of this work. \n3. The main result that the Fast Equilibrium conjecture holds without the assumptions of the unique basin and that the set of minima is an analytic set is significant, and it contributes well to the theoretical understanding of the effects of normalization layers. \n4. While I did not check all the proof details, I followed the proof at a high level.", "weaknesses": "It could be good to add some theoretical reasonings (in addition to being natural) about why Assumption 1.2 might be essential to prove the Fast Equilibrium conjecture.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GXk2iXlGw2", "reviewer": "Reviewer_LagX", "summary": "In this technical paper, the authors deal with fast convergence of networks with normalized steps to the solution.  The authors successfully prove their result and many insights are given on practical training of these networks in the experiments section.", "questions": "1) This is something I am pretty sure I could solve on my own with a bit of thinking, but I guess it hints to lack of clarity in some parts of the intro: I find a bit of contradiction between the sentences (1, abstract) \"scale-invariant normalized network, mixes to an equilibrium in\n$O(1/\\eta\\lambda)$ steps, as opposed to classical $e^{O(\\eta^{−1})}$ mixing time\" and (2, intro) \"When normalization is used,  the effective learning rate for the renormalized parameter vector will stabilize around $O((\\eta\\lambda^{-1/2} )$ and in consequence $e^{O((\\eta\\lambda)^{-1}}$ is replaced with $e^{O((\\eta\\lambda)^{-1/2}}$\". I think this is a bit unclear, can you please explain? \n\n2) From my SDE knowledge, I always understood that convergence to the stationary distribution is dominated by the drift. What is the convergence rate -- to a local minimum -- if you drop the noise term? I think providing an analysis of this setting will certainly help the reader understand the proof in a simplified setting. \n\n3) The SDE you study certainly is a model for SGD in the setting you study. However, I am a bit worried about the noise structure. Is there some guarantee in previous literature that constant gradient noise in the not-normalized setting translates to the noise projection structure you study in the normalized case? I am talking about formula 4-5 in comparison to the discrete update.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "This is a very technical paper, very mathematical, and centered around proving one conjecture in the literature. I like this, and I also enjoyed reading the paper: it is very well written and has pleasant notation. What I especially liked is the clear statement of the assumption and remarks 1.4 and 1.5. The authors also showed they mastered the subject with many useful citations and discussions around existing results. Though I did not unfortunately find the time to go through the proofs, the narrative and the impeccable discussion in the main paper leave no doubts on correctness. On a technical level, I was always curious about projections of the Brownian motion on the sphere in the context of Langevin dynamics, so I will for sure get back to this paper in the future to find more details on this. I also have a question, which you can find in the proper section below.\n\nI placed the contribution as fair since I think the result does not explain the fast convergence of normalized networks compared to not-normalized (motivation in the abstract). Please comment on this if you think your result motivates this, I am happy to revise!", "weaknesses": "I guess one obvious question is \"can you have more experiments\". I think this might be silly in this paper, but maybe there is something you can do to reach people outside the very technical domain. One thing I think would be useful is to illustrate the rates known in the literature and the conjecture - maybe with evidence from some datasets and some networks. I think you can attract the interest of many people if you have a headline figure showing the speed of convergence to the stationary distribution and exactly the rate you prove.\nA thing I found a bit confusing is going back-and-forward on between $O(1/\\eta)$ and $e^{O(\\eta^{-1})}$ results. That got me a bit thinking and I have a question (below).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ea6YD1xUJQ", "reviewer": "Reviewer_pBQz", "summary": "To understand the behaviors of normalization in deep learning, Li et al. (2020) proposes the Fast Equilibrium conjecture: the scale-invariant normalized network, when trained by SGD with $\\eta$ learning rate and $\\lambda$ weight decay, mixes to an equilibrium in $\\tilde{O}(\\frac{1}{\\eta \\lambda})$ steps, as opposed to classical $e^{O((\\eta \\lambda)^{-1})}$ mixing time. Recent works by Wang & Wang (2022) and Li et al. (2022c) further proved this conjecture under different sets of assumptions.\n\nThis paper instead proves the fast equilibrium conjecture in full generality by removing the non-generic assumptions of Wang & Wang (2022) and Li et al. (2022c) that the minima are isolated, that the region near minima forms a unique basin, and that the set of minima is an analytic set. Their main technical contribution is to show that with probability close to 1, in exponential time trajectories will not escape the attracting basin containing its initial position.", "questions": "I have a question about the classical mixing time: You use $e^{O(\\eta^{-1})}$ in the abstract but use $e^{O((\\eta \\lambda)^{-1})}$ in the Introduction. I feel $e^{O((\\eta \\lambda)^{-1})}$ looks more natural to me. Could you clarify this?\n\nSimilarly, $\\tilde{O}(1/\\eta\\lambda)$ in the abstract may need to be changed as $\\tilde{O}(\\frac{1}{\\eta \\lambda})$.", "limitations": "", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Originality:** 1) They first analyze the generality of assumptions used in existing work and then successfully remove non-generic assumptions, which is very important to reduce the gap between the theory and the experiments 2) They use Arnold-Kliemann's condition instead of Kliemann's condition to remove the analyticity assumption 3) They make use of large deviation principle of Dembo-Zeitouni in Dembo & Zeitouni (2010, Chapter 5) instead of the Freidlin-Wentzell’s original theory to show that with very high probability, the trajectory will not escape from the basin in exponential time to remove the unique basin assumption\n\n**Quality:** I'm not familiar with the theoretical techniques but I feel it has good quality.\n\n**Clarity:** It looks quite clear in most cases but needs some modifications for some tiny writing errors.\n\n**Significance:** Okay but not very significant for the following two reasons: 1) It relaxes the assumptions for an existing conjecture instead of discovering some new phenomena, 2) The fast equilibrium conjecture is not as important as many other things in deep learning, such as the generalization ability of deep neural networks, the puzzles in large language models, the interplay among the model, the algorithm, and the data, etc.", "weaknesses": "1) They only conduct experiments on MNIST which is almost a linearly-separable dataset, which is not good for deep learning analysis. I suggest the authors conduct the experiments on other more difficult datasets, such as CIFAR-10.\n2) There are some tiny errors in the paper, so I'd suggest the authors to proofread their paper more carefully. \n- They sometimes use an uncommon citation format for references in some places. For example,\n\"The works by Bovier et al. and Shi et al. (Bovier et al., 2004; Shi et al., 2020)\" may be changed as \"Bovier et al. (2004) and Shi et al. (2020)\"; \"Li et al. made certain assumptions in (Li et al., 2022c)\" may be changed as \"Li et al. (2022c) made certain assumptions\".\n- Similarly, they seem use wrong references in some places. For example,\n\"We now stop assuming Assumption 1.3.(ii) and decompose\" => I feel it's Assumption 1.3. (i) instead of Assumption 1.3 (i) since they are talking about removing the unique basin assumption (Assumption 1.3 (i));\n\"Recall that (Li et al., 2022c)) also assumes Assumption 1.3.(i), but that can be dropped by the discussion in Chapter 3 above.\" => I feel it's Assumption 1.3 (ii) instead of Assumption 1.3 (i) since Chapter 3 discusses removing the analyticity assumption (Assumption 1.3. (ii)); \n\"Figure 5 shows that V11 and V22 stabilizes near similar but different values,\" => I think they mean Figure 1 instead of Figure 5 here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)