Large Stepsize Gradient Descent for Non-Homogeneous Two-Layer Networks: Margin Improvement and Fast Optimization
Authors: Yuhang Cai
Jingfeng Wu
Song Mei
Michael Lindsey
Peter Bartlett
Venue: NeurIPS 2024 poster (2024)
Keywords: Deep Learning Theory
Large stepsize
Optimization
Implicit Bias
Primary area: optimization_for_deep_networks

Abstract: The typical training of neural networks using large stepsize gradient descent (GD) under the logistic loss often involves two distinct phases, where the empirical risk oscillates in the first phase but decreases monotonically in the second phase. We investigate this phenomenon in two-layer networks that satisfy a near-homogeneity condition. We show that the second phase begins once the empirical risk falls below a certain threshold, dependent on the stepsize. Additionally, we show that the normalized margin grows nearly monotonically in the second phase, demonstrating an implicit bias of GD in training non-homogeneous predictors. If the dataset is linearly separable and the derivative of the activation function is bounded away from zero, we show that the average empirical risk decreases, implying that the first phase must stop in finite steps. Finally, we demonstrate that by choosing a suitably large stepsize, GD that undergoes this phase transition is more efficient than GD that monotonically decreases the risk. Our analysis applies to networks of any width, beyond the well-known neural tangent kernel and mean-field regimes.

Large Stepsize Gradient Descent for Non-Homogeneous Two-Layer Networks: Margin Improvement and Fast Optimization
Yuhang Cai Jingfeng Wu² Song Mei
Michael Lindsey 1.2
Peter L. Bartlett 1.3
UC Berkeley "Lawrence Berkeley National Laboratory
GGoolle DeepMind
{willcai, uuujf, songmei, lindsey, peter Jeberkeley. edu
Abstract
The typical training neural networks using large stepsize gradient descent (GD) under the logistic loss often involves two distinct phases, where the empirical risk oscillates in the first phase but decreases monotonically the second phase. We investigate this phenomenon in two-layer networks that satisfy near-homogencity condition. We show that the second phase begins once the empirical risk falls below certain threshold, dependent the stepsize. Additionally, normalized margin grows nearly monotonically in the second phase, demonstrating implicit bias of GD in training non-homogeneous predictors. If the dataset is linearly separable and the derivative of the activation function is bounded away from zero, show that the average empirical risk decreases, implying that the first phase must stop in finite steps. Finally, we demonstrate that by choosing a suitably large stepsize, GD that undergoes this phase transition is more efficient than GD that monotonically decreases the risk. Our analysis applies to networks of any width, beyond the well-known neural tangent kernel and mean-field regimes.
show that the
Introduction
Neural networks are mostly optimized by gradient descent (GD) or its variants. Understanding the behavior of GD is one of the key challenges in deep learning theory. However, there is a nonnegligible discrepancy between the GD setups in theory and in practice. In theory, GD is mostly analyzed with relatively small stepsizes such that its dynamics are close to the continuous gradient flow dynamics, although few exceptions will be discussed later. However, in practice, GD is often used with relatively large stepsize, with behaviors significantly deviating from that of small stepsize GD or gradient flow. Specifically, notice that small stepsize GD (hence also gradient flow) induces monotonically decreasing empirical risk, but in practice, good optimization and generalization performance is usually achieved when the stepsize is large and the empirical risk oscillates [see Wu and Ma, 2018, Cohen et al., 2020, for example]. Therefore,
unclear which of the theoretical insights drawn from analyzing small stepsize GD apply to large stepsize GD used practically.
The behavior of small stepsize GD relatively well understood. For instance, classical optimization theory suggests that GD minimizes convex and L-smooth functions if the stepsize is well below 2/ with a convergence rate of ((1/((()) where t the number of steps [Nesterov, 2018]. More recently, Soudry et al. [2018]. Ji and Telgarsky [2018] show an implicit bias of small stepsize GD in logistic regression with separable data, where the direction of the GD iterates converges to the max-margi direction. Subsequent works extend their implicit bias theory from linear model to homogenous networks [Lyu and Li, 2020, Chizat and Bach, 2020, Ji and Telgarsky, 2020]. These theoretical results all assume the stepsize of GD is small (and even infinitesimal) such that the empirical risk decreases monotonically and, therefore cannot be directly applied to large stepsize GD used in practice.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

10
000007
D3
0- 0.00
000003
Komssked morgnn
Taiing boc.
0.0003
0.09
0.0009
607
0.2 0.003
000003
0.0065
12.3 1he 373
Eerations
Peraatoos
Berations
(a) Empirical risk
(b) Normalized margin
(c) Test accuracy
Figure The behavior of (GD) for optimizing a non-homogenous four-layer MLP with GELU activation function on subset of CIFAR-10 dataset. We randomly sample 6, 000 data with la- bels "airplane" and "automobile" from CIFAR- dataset. The normalized margin is defined as min;gim] 34 (w;; x4)//1)/ll xx)/|[w,|l| which is close to (3). The blue curves correspond GD with large stepsize 7j = 0.2, where the empirical risk oscillates the first phase but decreases monotonically in the second phase. The orange curves correspond to GD with small stepsize 0.005, where the empirical risk decreases monotomically. Furthermore, Figure 1(b) suggests the normalized margins of both two curves increase and converge in the stable phases. Finally. Figure I(c) suggests that large stepsize achieves better test accuracy, consistent with larger-scale learning experiment [Hoffer et al., 2017, Goyal al., 2017]. More details be found in Section 5.
More recently, large stepsize GD that induces oscillatory risk has been analyzed in simplified setups [see Ahn al., 2023, Zhu al., 2022, Kreisler al., 2023, Chen and Bruna, 2023, Wang al., 2022a, Wu et al., 2023, 2024, for an incomplete list of references]. In particular, in logistic regression with linearly separable data, Wu al. [2023] showed that the implicit bias of GD (that maximizes the margin) holds not only for small stepsizes [Soudry et al., 2018. Ji and Telgarsky, 2018] 8] but also for an arbitrarily large stepsize. In the same problem, Wu al. [2024] further showed that large stepsize GD that undergoes risk oscillation can achieve (L/t²²) empirical risk. whereas small stepsize GD that monotonically decreases the empirical risk must suffer from (2(1tt) empirical risk. Nonetheless, these theories of large stepsize GD are limited relatively simple setups such as linear models. The theory of large stepsize GD for nonlinear networks is underdeveloped.
This work fills the gap by providing an analysis of large stepsize GD for nonlinear networks. In the following, we set up our problem formally and summarize our contributions.
Setup. Consider binary classification dataset (x,, 34) where x, E Rd is feature vector and V: E {+1} is a binary label. For simplicity, we assume ||x; < 1 for all throughout the paper. For a predictor f, the empirical risk under logistic loss is defined
L(w) = Σ ((yaf(w;x))) 2(1) = log(1 + e
(I)
Here, the predictor /(w: :) Rd + is parameterized by trainable parameters w and is assumed be continuously differentiable with respect to w. The predictor is initialized from wo and then trained by gradient descent (GD) with a constant stepsize 7j > 0, that is,
WH1 == w. iwWW,.,) t 0.
(GD)
We are interested in nonlinear predictor f and a large stepsize rj. A notable example in our theory is two-layer networks with Lipschitz, smooth, and nearly homogenous activations (see (2)). Note that minimizing T.(w) is non-convex problem in general.
Observation. Empirically, large stepsize GD often undergoes a phase transition, where the empirical risk defined in (1) oscillates in the first phase but decreases monotonically in the second phase (see empirical evidence Appendix in [Cohen et al., 2020] and formal proof in [Wu et al., 2024] for linear predictors). This is illustrated in Figure 1. We follow Wu et al. [2024] and call the two phases the edge of stability (FoS) phase [name coined by Cohen et al., 2020] and the stable phase, respectively.

---

Contributions. We prove the following results for large stepsize GD for training nonlinear predic- tors under logistic loss.
1. For Lipschitz and smooth predictor / trained by GD with stepsize vj, we show that long the empirical risk is below a threshold depending on ij, GD monotonically decreases the empirical risk (see Theorem 2.2). This result extends the stable phase result in Wu et al. [2024] from linear predictors to nonlinear predictors, demonstrating the generality of the existence of a stable phase.
Assuming that GD enters the stable phase, if in addition the preditor has a bounded ho- mogenous error (see Assumption IC), we show that the normalized margin induced by GD, min; (w;; x nearly monotonically increases (see Theorem 2.2). To the best of our knowledge, this is the first characterization of implicit bias of GD for non-homogenous predictors. In particular, our theory covers two-layer networks with commonly used activations functions (which are often non-homogenous) that cannot be covered by existing results [Lyu and Li, 2020, Ji and Telgarsky, 2020, Chizat and Bach, 2020].
Under additional technical assumptions (the dataset is linearly separable and the derivative of the activation function is bounded away from zero), we show that the initial EoS phase must stop in (1)) steps and GD transits to the stable phase afterwards. Furthermore, by choosing a suitably large stepsize, GD achieves a (1((t) empirical risk after t steps. In comparison, GD that converges monotonically incurs an (11/t) risk. This result indicates optimization benefit of using large stepsize and generalizes the results in [Wu et al., 2024] from linear predictors to neural networks.
2 Stable Phase and Margin Improvement
In this section, we present our results for the stable phase of large stepsize GD in training nonlinear predictors. Specifically, our results apply nonlinear predictors that are Lipschitz, smooth, and nearly homogeneous, as described by the following assumption.
Assumption (Model conditions). Consider a predictor (W;;))) where x, is one of the feature vectors in the training set.
A. Lipschitznes. Assume there exists > 0 such that for every w, sup, ///.f/::;x/// S
B. Smoothness. Assume there exists 3 > 0 such that for all w, V,
[V;(w; (v: S Bl[w vll, i i= I....,n
C. Near homogeneity. Assume there exists such that for every w, If(w: W (w; w)l K. i= ==
Assumptions 1A and IB are commonly used conditions in the optimization literature. Note that Assumption 1B implies continuous differentiabilityy thus ruling out networks with ReLU activation function. The continuous differentiabilityy is only used our current stable phase analysis. We conjecture it can be relaxed using subdifferentiability [Lyu and Li, 2020] for allowing ReL.U networks.
If Æ then Assumption IC requires the predictor be exactly 1-homogenous. Our Assumption IC allows the predictor to have a bounded homogenous error. It is clear that linear predictors f(w; x) :=
|| S 1, B = O, and A = O. Another notable example is w T x satisfy Assumption with = sup; |xx|| two-layer networks given by
f(w: x) = m Σ ,x (w))),
w E ER, j= 1-
(2)
where we assume a, {+1} are fixed and w E Rmd the stack of (wi)) )- are the trainable parameters. We define two-layer networks with the mean-field scaling [Song et al., 2018, Chizat and Bach, 2020, Chen et al., 2022, Suzuki et al., 2023]. However, our results hold for any width. The effect of rescaling the model will be discussed in Section 4. The following example shows that Assumption covers two-layer networks with many commonly used activations ((-). The proof provided in Appendix F.I.
3

---

Example 2.1 (Two-layer networks), Two-layer networks defined in (2) with the following activation
functions satisfy Assumption I with the described constants:
GELU. (xx) = [e + (x/v2)) with = 6-1/2/12T, 3 2/m, and p (V2T
e -1/2)/
Softplus. ((x)) = log(1 + e*) with &. = log 3 = 1/m, and p = 1/mm.
SiLU. Φ(z) 22((1 ") with K 1, 5 4/m, and p 2/ Vm.
Huberized ReLU [Chatterji et al., 2021]. For a fixed h >
x < 0,
<(x) =
x²
0 x S h,
2h
h,
with h/2. 3 1/(mm), and p 1/m
Margin for nearly homogenous predictors. For nearly homogenous predictor f(w;-) (see Assumption IC), we define its normalized margin (or margin for simplicity)
7(w) minge[n) 3r (w::
|w|
large normalized margin (w) guarantees the prediction of each sample is away from the decision boundary. The normalized margin (3) is introduced by Lyu and Li [2020] for homogenous predictors. However, we show that the same notion is also well-defined for non-home homogenous predictors that satisfy Assumption IC. The next theorem gives sufficient conditions for large stepsize GD enter the stable phase in training non-homogenous predictors and characterizes the increase of the normalized margin. The proof of Theorem 2.2 is deferred Appendix A.
Theorem 2.2 (Stable phase and margin improvement). Consider (GD) with stepsize ij on predictor f(w; x) that satisfies Assumptions 1A and IB. If there exists such that
L(we) S 7(2p² + B)
then GD is in the stable phase for 1 2 r, that is, (I(w)))rzr decreases monotonically. additionally the predictor satisfies Assumplion 1C and there exists 8 2 0 such that
1
L(W,) S min
(5)
ex+22m
7(40² + 23)
we have the following for trrt2 2
Risk convergence. L(wi)= = (11/t).
Parameter increase. ||W:+1|| 2 ||well and ||wlll = O(log(!))
Margin improvement. There exists a modified margin function (w) such that
(w:) is increasing and bounded.
7 (w:) is a multiplicative approximiator of (we), that is, there exists c > such that
(w)) S ((w)) < [ log(L/L(w:))
(we)), t
As direct consequence, lim 7w:)) lim (wt).
Theorem 2.2 shows that for an arbitrarily large stepsize 1j. GD must enter the stable phase the empirical risk falls below threshold depending on 7 given by (4). Furthermore, for nearly homogenous predictors, Theorem 2.2 shows that under a stronger risk threshold condition (5), the risk must converge at ((1/L) rate and that the normalized margin nearly monotonically increases. This demonstrates an implicit bias of GD, even when used with a large stepsize and the trained predictor is non-homoge.ous.

---

Limitations. The stable phase conditions in Theorem 2.2 require GD to enter sublevel set of the empirical risk. However, such a sublevel set might be empty. For instance, let f fw::x) be a two-layer network (2) with sigmoid activation. Notice that the predictor is uniformly bounded, [s(w; x) S so we have
L(w) Σ log(1 + e -uf((wo))) 2 log(1 + te
On the other hand, we can also verify that Assumption IC is satisfied by f(w; x) with K = 1 but no smaller K. Therefore (5) cannot be satisfied. In general, the sublevel set given by the right-hand side of Assumption 1C is non-empty
there exists a unit vector v such that min y (Av; - D0 as oo.
The above condition requires the data can be separated arbitrarily well by some predictor within the hypothesis class. This condition is general and covers (sufficiently large) two-layer networks (2) with many commonly used activations such as GeLU and SiLU. Moreover, although two-layer networks with sigmoid activation violate this condition, they be modified by adding leakage to the sigmoid to satisfy the condition. Furthermore, this condition can be satisfied for some nonlinear problems like XOR (or k-parity problems) since they can be realized by two-layer networks.
In the next section, we will provide sufficient conditions such that large stepsize GD will enter the stable phases characterized by (4) or (5).
Comparisons to existing works. Our Theorem 2.2 makes several important extensions compared existing results [Wu et al., 2024, Lyu and Li, 2020, Chizat and Bach, 2020, Ji and Telgarsky, 2020]. First, Theorem 2.2 suggests that the stable phase happens for general nonlinear predictors such as two-layer networks, while the work by Wu et al. [2024] only studied the stable phase for linear predictors. Second, the margin improvement only known for small (and even infinitesimal) stepsize GD and homogenous predictors [Lyu and Li, 2020, Chizat and Bach, 2020, Ji and Telgarsky, 2020], while we extend this to non-homogenous networks. To the best of our knowledge, Theorem 2.2 is the first implicit bias result covering large stepsize GD and non-hoomgeeous predictors.
From a technical perspective, our proof uses techniques introduced by Lyu and Li [2020] for analyzing homogenous predictors. Our main innovation is the construction of new auxiliary margin functions that can deal with errors caused by large stepsize and non-homogeneity. More details are discussed in Appendix A.3.
3 Edge of Stability Phase
Our stable phase results in Theorem 2.2 require the risk to be below certain threshold (see (4) and (5)). this section, we show that the risk can indeed be below the required threshold, even when GD is used with large stepsize. Recall that minimizing the empirical risk with a nonlinear predictor is non-convex, therefore solving it by GD is hard general. We make additional technical assumptions to conquer the challenges caused by non-convexity. We conjecture that these technical assumptions are not necessary and can be relaxed.
We focus on two-layer networks (2). We make the following assumptions on the activation function.
Assumption 2 (Activation function conditions). In the two-layer network (2), let the activation function o R R be continuously differentiable. Moreover,
A. Derivative condition. Assume there exists 0 < <a < 1 such that a S [0'(2) 1
Smoothness. Assume there exists B > such that for all x, y E R. (0'(x) - ('(v) S B|x y/-
Near homogeneity. Assume there exists
such that for every 2 R. (Ø(2) (z)22) S K.
Recall that sup; ||x;| < 1. One can then check by direct computation that, under Assumption 2, two-layer networks (2) satisfy Assumption 1 with p = 1/vm, 3/m, and k.
Assumptions 2B and 2C cover many commonly used activation functions. In Assumption 2A, we assume (ø((z) 1. This is just for the simplicity of presentation and our results can be easily
5

---

generalized to allow [0(22) I C for a constant C > 0. The other condition in Assumption 2A, (2)) 2 a, however, non non-trivial. This condition is widely used literature [see Brutzkus et al., 2018, Frei et al., 2021, and references thereafter] to facilitate GD analysis. Technically, this condition guarantees that each neuron in the two-layer network (2) will always receive a non-trivial gradient in the GD update; otherwise, neurons may be frozen during the GD update. Furthermore, commonly used activation functions can be combined with an identity map to satisfy Assumption 2A. This is formalized in the following example. The proof is provided in Appendix F.2.
Example 3.1 (Leaky activation functions). Fix 0.5 S ccl
Let o be GELU, Softplus, SilU in Example 2.1, then its modification oι) cr+ (1-c)r) satisfies Assumption 2 with L. 0.25, and 3 = 1. In particular, the modification of softplus can be viewed as a smoothed leaky ReLU.
Let o be the Huberized ReLU in Example 2.1, then its modification Ø(z) == cr (1- c)/4 o(x)
satisfies Assumption 2 with h/2, 0.5, and / 1/4h.
The "leaky" tanh, x(x) = cz - (1 - c) tanh(x), and the "leaky" sigmoid, x(x) = cx c/(1+e both satisfy Assumption 2 with x = 1, 0.5 and 3 =.
For the technical difficulty of non-convex optimization, we also need to assume a linearly separable dataset to conduct our BoS phase analysis.
Assumption 3 (Linear separability). Assume there is margin y > 0 and a unit vector w. suschh that y:x; w. 2 y for every ,n.
Assumption 3 serves as sufficient condition for two-layer neural networks, regardless of width, to reach the initial bound the stable phase under large stepsizes. We remark that our stable phase results do not need this assumption.
The following theorem shows that when GD is used with large stepsizes, the average risk must decrease even though the risk may oscillate locally.
Theorem 3.2 (The EoS phase for two-layer networks). Under Assumption 3, consider (GD) on two-layer networks (2) that satisfy Assumptions 2A and 2C. Denote the stepsize by == mn, where m is the network width and can be arbitrarily large. Then for every t > 0, we have
1 8 log (7 nt? + Sn. 8n²//² +
|wol? must
o log (nt) + n?
L(WA) <
+
2tt
rtt
A=0
Theorem 3.2 suggests that the average risk training two-layer networks decreases even when GD used with large stepsize. Consequently, the risk thresholds (4) and (5) for GD enter the stable phase must be satisfied after a finite number of steps. This will be discussed in depth in the next section.
Compared to Theorem in [Wu et al., 2024], Theorem 3.2 extends their EoS phase bound from linear predictors to two-layer networks.
4 Phase Transition and Fast Optimization
For two-layer networks trained by large stepsize GD, Theorem 3.2 shows that the average risk must decrease over time. Combining this with Theorem 2.2, GD must enter the stable phase in finite steps, and the loss must converge while the normalized margin must improve.
However, direct application of Theorem 3.2 only leads to suboptimal bound on the phase transition time. Motivated by Wu et al. [2024], we establish the following sharp bound on the phase transition time by tracking the gradient potential (see Lemma C.3). The proof is deferred to Appendix C.
Theorem 4.1 (Phase transition and stable phase for two-layer networks). Under Assumption 3, consider (GD) on two-layer networks (2) that satisfy Assumption 2. Clearly, the two-layer networks also satisfy Assumption I with p = 1/ /m, 3 = 3/m, and & = k. Denote the stepsize by ñ == mn. where is the network width and 17 > 0 can be arbitrarily large.

---

Phase transition time. There exists s S T such that (5) in Theorem 2.2 holds, where
7 128(1 + 4k)
max C177, C,n, (2)) G1 log
(2)) + cin
(22)) cin)
| v
α
"
where C1 == 4e and eg == (8 + 48). Therefore (GD) is in the stable from s onwards.
Explicit risk bound the stable phase. We have [[(w)))2e monotonically decreuses and
2 L(w,) a2 s)
Theorems 2.2, 3.2 and 4.1 together characterize the behaviors of large stepsize GD in training two- layer networks. Specifically, large stepsize GD may induce an oscillatory risk in the beginning; but the averaged empirical risk must decrease (Theorem 3.2). After the empirical risk falls below certain sepsize-deeendeen threshold, GD enters the stable phase, where the risk decreases monotonically (Theorem 4.1). Finally, the normalized margin (3) induced by GD increases nearly monotonically as GD stays in the stable phase (Theorem 2.2).
Our intuition behind the phase transition phenomenon is as follows. The initial EoS phase occurs when gGD oscillates within a steep valley, transitioning to a stable phase once it navigates into flatter valley. We believe this insight generalizes broader nonlinear models. Moreover, our theory of large step sizes aligns with the celebrated flat minima intuition [Keskar et al., 2016].
Fast optimization. Our bounds for two-layer networks are comparable to those for linear predictors shown by Wu et al. [2024]] Specifically, when used with a larger stepsize, GD achieves faster optimization in the stable phase but stays longer in the EoS phase. Choosing a suitably large stepsize that balances the steps EoS and stable phases, we obtain an accelerated empirical risk in the following corollary. The proof is included in Appendix C.2.
Corollary 4.2 (Acceleration of large stepsize). Under the same setup as in Theorem 4.1, consider (GD) with a given budget of T steps such that
256(1 + 4k) T2
max { (in, 4c2,
20z||wol|
a27²
vm
where = 4+2 and 02 = (8 + 43) are in Theorem 4.1. Then for stepsize क == 1871 where
a²²² T,
= 256(1 + 4x)c2
we have S 2/2 and
(048(1- + 4x:002
L(Wr)
O
at74
T2
T
Theorem 4.1 and Corollary 4.2 extend Theorem I and Corollary in Wu et al. [2024] from linear predictors to two-layer networks. Another notable difference is that we obtain sharper stable phase bound (and thus better acceleration bound) compared theirs, where we remove logarithmic factor through more careful analysis.
Corollary 4.2 suggests an accelerated risk bound of O(1/T²) by choosing large stepsize that balances EosS and stable phases. We also show the following lower bound, showing that such accel- eration is impossible if (GD) does not enter the EoS phase. The proof is included in Appendix C.3.
Theorem 4.3 (Lower bound in the classical regime). Consider (GD) with initialization wo 0 and stepsize îj > O for a two-layer network (2) satisfying Assumption 2. Suppose the training set given by
x1
(66 -1-12)))
X2 (11:-v1-22700 V1 U2 1, 0 0.1.
is clear that (x,, M1) satisfy Assumption 3. If (L (we)) is non-increasing. then
L ( 2 co/t, 1
where co > is function of (a, d, x,, X2, K,K, but is independent of and îj.
7

---

Effect of model rescaling. We conclude this section by discussing the impact of rescaling the model. Specifically, we replace the two-layer network the mean-field scaling (2) by the following
f(w; x) = b.
m E aj0(x Tw())
1
m
and evaluate the impact of the scaling factor b on results. By choosing the optimal stepsize that balances the EoS and stable phases as in Corollary 4.2, we optimize the risk bound obtained by GD with fixed budget of T steps and get the following bound. Detailed derivations deferred to Appendix D.
(L(1/29) L(Wr) (0(
if 1, if 1.
This suggests that long b 2 1, get the acceleration effect. In particular, the mean-field scaling b [Song et al., 2018, Chizat and Bach, 2020] and the neural tangent kernel (NTK) scaling b = mm [Du al., 2018, Jacot et al., 2018] give the same acceleration effect. An NTK analysis of large stepsize is included in [Wu et al., 2024] and their conclusion is consistent with ours. Finally, we remark that our analysis holds for any width m and uses techniques different from the mean-field or NTK methods. However, our acceleration analysis only allows linearly separable datasets.
5 Experiments
We conduct three sets of experiments to validate theoretical insights. In the first set, subset of the CIFAR-10 dataset [Krizhevsky et al., 2009], which includes 6,000 randomly selected samples from the "airplane" and "automobile" classes. Our model is multilayer perceptron (MLP) with four trainable layers and GELU activation functions, with a hidden dimension of 200 for each hidden layer. The MLP trained using gradient descent with random initialization, described in (GD). The results are shown in Figures ((a) to (CC).
In the second set of experiments, consider an XOR dataset consisting of four samples:
X = (-1,-1))
x2 = (1,1), 1; X3 (1, (1), V/ 1; x4 (-1,1), V4 1.
The above XOR dataset is not linearly separable. We test (GD) with different stepsizes on a two-layer network (2) with the leaky softplus activation (see Example 3.1 with = 0.5). The network width is m 20. The initialization random. The results are presented in Figures 2(a) to 2(c).
In the third set of experiments, we consider the same task as in the first set of experiments, but we test (GD) with different stepsizes two-layer network (2) with the softplus activation. The network width is m 40. The initialization is random. The results are presented in Figures 2(d) to 2(1).
Margin improvement. Figures 1(b), 2(c) and 2(f) show that the normalized margin nearly mono- tonically increases once gradient descent (GD) enters the stable phase, regardless of step size. This observation aligns with our theoretical findings in Theorem 2.2.
Fast optimization. From Figures (a), 2(a) and 2(d), we observe that after GD enters the stable phase, larger stepsize consistently leads to smaller empirical risk compared the smaller stepsizes, which is consistent with our Theorem 4.1 I and Corollary 4.2. Besides, Figures 2(b) and 2(e) suggest that, asymptotically, GD converges rate of (ii(()t)) O(1(m)() (The width of networks is fixed), which verifies the sharpness of our stable phase bound in Theorem 4.1.
Margin of individual neurons. It is important to note that while the normalized margin behaves as expected, the margin for individual neurons may not increase and can remain negative, even when the dataset is linearly separable. detailed example illustrating this provided in Appendix E.
6 Related Works
In this section, we discuss related papers.
8

---

0310
12
0.13
Liwel -0
0.33
0.12
0.65
0.63
0085
(a) Empirical risk, XOR.
(b) Asymptotic rate, XOR.
(c) Normalized margin, XOR.
05
127
0.5
h- 025
1= 0.25
1000
E43
-14--
-122
127
1305
-103
200
125
(d) Empirical risk, CIFAR-10.
(e) Asymptotic rate, 10FAR-10.
(f) Normalized margin, CIFAR-10.
Figure 2: Behavior of (GD) for two-layer networks (2) with leaky softplus activation function (see Example 3.1 with e 0.5). We consider an XOR dataset and subset of CIFAR- 10 dataset. both cases, we observe that (1) GD with large stepsize achieves a faster optimization compared to GD with small stepsize, (2) the asymptotic convergence rate of the empirical risk is ((1((7()), and (3) in the stable phase, the normalized margin (nearly) monotonically increases. These observations are consistent with our theoretical understanding of large stepsize GD. More details about the experiments are explained in Section 5.
Small stepsize and implicit bias. For logistic regression on linearly separable data, Soudry et al. [2018], Ji and Telgarsky [2018] showed that the direction of small stepsize GD converges to the max-margin solution. Their results were later extended by Gunasekar et al. [2017, 2018], Nacson et al. [2019c,a,b], Ji et al. [2021], Lyu and Li [2020], Ji and Telgarsky [2020], Chizat and Bach [2020], Chatterji et al. [2021], Kunin et al. [2022] to other algorithms and non-linear models. However, in all of their analysis, the stepsize of GD needs to be small such that the empirical risk decreases monotonically. In contrast, our focus is GD with large stepsize that induces non-monotonic risk.
Two papers [Nacson et al., 2019a, Kunin et al., 2022] studied margin maximization theory for special form of non-homogenous models. Specifically, when viewed terms of different subsets of the trainable parameters, the model is homogencous, although the order of homogeneity may vary. Compared to their setting, our non-homogenous models only require a bounded homogenous error (see Assumption IC). Therefore, our theory cover two-layer networks (2) with non-homogeneous activations such as GELU and SiLU that cannot be covered by [Nacson et al., 2019a, Kunin et al., 2022].
Large stepsize and EoS. In practice, large stepsizes are often preferred when using GD to train neural networks achieve effective optimization and generalization performance [see Wu and Ma, 2018, Cohen et al., 2020, Barrett and Dherin, 2020, and references therein]. In such scenarios, the empirical risk often oscillates in the beginning. This phenomenon is named edge of stability (EoS) by Cohen et al. [2020]. The theory of EoS is mainly studied in relatively simplified cases such as one- or awo-dimensional functions [Zhu al., 2022, Chen and Bruna, 2023, Ahn al., 2022, Kreisler et al., 2023, Wang et al., 2023], linear model [Wu et al., 2023, 2024], matrix factorization [Wang et al., 2022a, Chen and Bruna, 2023], scale-invariant networks [Lyu et al., 2022], línear networks under MSE loss [Ren et al., 2024, Andriushchenko et al., 2023], for an incomplete list of references. Compared to them, we focus on more practical setup of training two-layer non-linear networks with large stepsize GD. There are some general theories of EoS subject to subtle assumptions [for

---

example, Kong and Tao, 2020, Ahn et al., 2022, Ma et al., 2022, Damian et al., 2022, Wang et al., 2022b, Lu et al., 2023], which are not directly comparable ours.
In what follows, we make a detailed discussion about papers that directly motivate our work yuu and Li, 2020, Ji and Telgarsky, 2020, Chatterji al., 2021, Wu al., 2024].
Comparison with Lyu and Li [2020], Ji and Telgarsky [2020]. Both results in [Lyu and Li, 2020, Ji and Telgarsky, 2020] focused on L-homogenous networks. Specifically, Lyu and Li [2020] showed that a modified version of normalized margin (see (3)) induced by GD with small stepsize (such that the risk decreases monotonically) increases, with limiting points of {we/ //el] converging to KKT points of mangin-maximizatiion problem. Under additional o-minimal conditions, Ji and Telgarsky [2020] showed that gradient flow converges in direction. Our work is different from theirs in two aspects. First, we allow GD with a large stepsize that may cause risk oscillation. Second, our theory covers non-homogenous predictors, which include two-layer networks with many commonly used activation functions beyond the scope of [Lyu and Li, 2020, Ji and Telgarsky, 2020]. Compared to Lyu and Li [2020], and Telgarsky [2020], we only show the improvement of the margin, and our theory is limited to nearly 1-homogenous predictors (Assumption 2C). It remains open to show directional convergence and to extend our near L-homogenity condition to a "near L-homogeneity" condition for general L.
Comparison with Chatterji et al. [2021]. The work by Chatterji et al. [2021] studies the conver- gence of GD in training deep networks under logistic loss. Their results are related to ours we both consider networks with nearly homogeneous activations and we both have a stable phase analysis (al- though this is explicitly mentioned in their paper). However, our results from theirs. Specifically, in our notation, they require the homogenous error ks (see Assumption 2C) to be smaller than (log(1/L(w, stable phase. Note that the margin when GD enters the stable phase could be arbitrarily small. In comparison, we only require the homogenous error be bounded by a constant. As a consequence, we can handle many commonly used activation functions (see Example 2.1) while they can only handle the Huberized ReLU with small Example 2.1. Moreover, they require the stepsize ij to be smaller than (x:/| [)) thus they only allow very small stepsize. In contrast, we allow rj to be arbitrarily large.
significantly different
w Il) = (W)())). where s is the time for GD to enter the
Comparison with Wu et al. [2023, 2024|. The works by Wu et al. [2023, 2024] directly motivate our paper. In particular, for logistic regression on linearly separable data, Wu et al. [2023] showed margin maximization of GD with large stepsize and Wu et al. [2024] showed fast optimization of GD with large stepsize. Our work can be viewed as an extension of [Wu et al., 2023, 2024] from linear predictors non-linear predictors such two-layer networks. Besides, our results for margin improvement and convergence within the stable phase (Theorem 2.2) hold for the general dataset, while their results strongly rely on the linear separability of the dataset.
7 Conclusion
We provide theory of large stepsize gradient descent (GD) for training non-homogeneous predictors such as two-layer networks using the logistic loss function. Our analysis explains the empirical observations: large stepsize GD often reveals two distinct phases in the training process, where the empirical risk oscillates in the beginning but decreases monotonically subsequently. We show that the phase transition happens because the average empirical risk decreases despite the risk oscillation. In addition, we show that large stepsize GD improves the normalized margin in the long run, which extends the existing implicit bias theory for homogenous predictors to non- non-homogenous predictors. Finally, we show that large stepsize GD, by entering the initial oscillatory phase, achieves acceleration when minimizing the empirical risk.
Acknowledgements
We thank Fabian Pedregosa for his suggestions on an early draft and Jason D. Lee and Kaifeng Lyu for their comments clarifying the applicability of our result to sigmoid networks. We gratefully acknowledge the support of the NSF for FODSI through grant DMS-2023505, of the NSF and the
10

---

Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMS-2031883 and #8 1814399, and of the ONR through MURI award N000142112431.
References
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. In International Conference Machine Learning, pages 247-257. PMLR, 2022.
Kwangjun Ahn, Sebastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via edge of stability. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023.
Maksym Andriushchenko, Aditya Vardhan Varre, Loucas Pillaud-Vivien, and Nicolas Flammarion. Sgd with large step sizes learns sparse features. In International Conference on Machine Learning, pages 903-925. PMLR, 2023.
David GT Barrett and Benoit Dherin. Implicit gradient regularization. arXiv preprint arXhv:2009.11162 2020.
Alon Brutzkus, Amir Globerson, Eran Malach, and Shai Shalev-Shwand Sgd learns parameterized networks that provably generalize on linearly separable data. In International Conference on Learning Representations, 2018.
Niladri S Chatterji, Philip M Long, and Peter Bartlett. When does gradient descent with logistic loss interpolate using deep networks with smoothed relu activations? In Conference on Learning Theory, pages 927-1027. PMLR, 2021.
Fan Chen, Zhenjie Ren, and Songbo Wang. Uniform-in-tinn -nttim propagation of chaos for mean field langevin dynamics. arXiv preprint arXiv:2212.03050, 2022.
Lei Chen and Joan Bruna. Beyond the edge of stability via two-step gradient updates. In International Conference on Machine Learning, pages 4330-4391. PMLR, 2023.
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on learning theory, pages 1305-1338. PMLR, 2020.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. International Conference on Learning Representations, 2020.
Alex Damian, Eshaan Nichani, and Jason D Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In The Eleventh International Conference on Learning Representa- tions, 2022.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations, 2018.
Spencer Frei, Yuan Cao, and Quanquan Gu. Provable generalization of sgd-trained neural networks of any width in the presence of adversarial label noise. In International Conference on Machine Learning, pages 3427-3438. PMLR, 2021.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in hour. arXiv preprint arXivv/706.0266 2017.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in Neural Information Processing Systems, 31, 2018.
11

---

Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: Closing the generaliza- tion gap large batch training of neural networks. Advances in Neural Information Processing Systems, 30, 2017.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems, 31, 2018.
Ziwei Ji and Matus Telgarsky. Risk and parameter convergence of logistic regression. arXiv preprint arXivv1803.073300 2018.
Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems, 33:17176-17188 2020.
Ziwei Ji, Nathan Srebro, and Matus Telgarsky. Fast margin maximization via dual acceleration. In International Conference on Machine Learning, pages 4860-4869. PMLR, 2021.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2016.
Lingkai Kong and Molei Tao. Stochasticity of deterministic gradient descent: Large learning rate for multiscale objective function. Advances in Neural Information Processing Systems, 33:2625-2638, 2020.
Itai Kreisler, Mor Shpigel Nacson, Daniel Soudry, and Yair Carmon. Gradient descent monotonically decreases the sharpness of gradient flow solutions in scalar networks and beyond. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 17684-17744, PMLR, 23-29 Jul 2023.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images, 2009.
Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maximum margin bias of quasi-homogeneous neural networks. In The Eleventh International Conference Learning Representations, 2022.
Miao Lu, Beining Wu, Xiaodong Yang, and Difan Zou. Benign oscillation of stochastic gradient descent with large learning rate. In The Twelfth International Conference on Learning Representa- tions, 2023.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2020.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction. Advances in Neural Information Processing Systems, 35: 14689-34770. 2022.
Chao Ma, Daniel Kunin, Lei Wu, and Lexing Ying. Beyond the quadratic approximation: The multiscale structure of neural network loss landscapes. arXiv preprint arXiv:2204.11326. 2022.
Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro. and Daniel Soudry. Lexicographic and depth-sensitive margins in homogeneous and non-homogeneous deep models. In International Conference on Machine Learning, pages 4683-4692. PMLR, 2019a.
Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3420-3428. PMLR, 2019b.
Mor Shpigel Nacson, Nathan Srebro, and Daniel Soudry. Stochastic gradient descent separable data: Exact convergence with a fixed learning rate. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3051-3059. PMLR, 2019c.
Yurii Nesterov. Lectures on Convex Optimization. Springer Publishing Company, Incorporated, 2nd edition, 2018. ISBN 13791.
12

---

Albert BJ Novikoff. On convergence proofs perceptrons. In Proceedings of the Symposium on the Mathematical Theory of Automata. New York, NY, 1962.
Yinuo Ren, Chao Ma, and Lexing Ying. Understanding the generalization benefits of late learning rate decay. In International Conference Artificial Intelligence and Statistics, pages 4465-4473.
PMLR, 2024.
Mei Song, Andrea Montanari, and P Nguyen. A mean field view of the landscape of two-layers neural networks. Proceedings of the National Academy of Sciences, 115(33),E7665-E767)) 2018.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1): 2822-2878, 2018.
Taiji Suzuki, Denny Wu, and Atsushi Nitanda. Convergence of mean-field langevin dynamies: Time and space discretization, stochastic gradient, and variance reduction. arXiv preprint 2rXiv:2306.07221,, 2023.
Yuging Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogeneity: Convergence and balancing effect. In International Conference Learning Representations, 2022a.
Yuging Wang, Zhenghao Xu, Tuo Zhao, and Molei Tao. Good regularity creates large learning rate implicit biases: edge of stability, balancing. and catapult. arXiv preprint arXiv:23 17087, 2023.
Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along GD trajectory: Progressive sharpening and edge of stability. Advances in Neural Information Processing Systems, 35:9983- 9994, 2022b.
Jingfeng Wu, Vladimir Braverman, and Jason D. Lee. Implicit bias of gradient descent for logistic regression the edge of stability. In Thirty-seventh Conference Neural Information Processing Systems, 2023.
Jingfeng Wu, Peter L Bartlett, Matus Telgarsky, and Bin Yu. Large stepsize gradient descent for logistic loss: Non-monotonicity of the loss improves optimization efficiency. Conference Learning Theory, 2024.
Lei Wu and Chao Ma. How SGD selects the global minima in over- wverppaamettrized learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, 31, 2018.
Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding odge-of-stability training dynamics with minimalist example. In The Eleventh International Conference on Learning Representations, 2022.
13

---

A Stable Phase Analysis
In this section, we will prove results for general smooth predictor f fw;; x) under the logistic loss in the stable phase. Before the proof, we introduce some notations here.
Notation.
We use the following notation to simplify the presentation.
(t) = yef (w,;x,),, 9min (t) = min (t(t).
L, L(w.), Pt //
Then, we have the following expression:
I(w.) Σ (()(1))-
Here, we give a summary of this section. The proofs are organized into parts.
In Appendix 1. we characterize the decrease of loss Lt.
In Appendix A.2, characterize the change of the parameter norm Pt-
In Appendix A.3, we show the convergence of the normalized margin 7(w.).
In Appendix A.4, we characterize the sharp rates of loss L, and parameter norm
In Appendix A.5, we give the proof of Theorem 2.2.
A.1 Decrease of the Loss
In this section, we will show that the loss L, decreases monotonically in the stable phase. To begin with, we introduce the following definition which is another characterization of f-smoothness.
Definition (Lincarization error). Given continuously differentiable function f E. and two points w, v E Rd, the linearization error of f(v) with respect to w is:
[[f](w, v) f(v) f(w) ff(w) (ow))((v--- -v-w)
For f-smooth function, standard convex optimization theory gives the following linearization error bound.
Fact A.1 (Linearization error of 3-smooth function). For a B-smooth function f Rd + R, we have
[[]](w, v) = f()) f(w) Vf(w) w) I 3
l|v pir-wwly
wll2. for every w and v.
We first show stable phase bound for general smooth and Lipschitz predictors. The following is an extension of Lemma 10 in [Wu et al., 2024]. Since we do not require ff to be twice differentiable, extra efforts are needed.
Lemma A.2 (Self-boundedness of logistic loss). For the logistic loss f(z) = log(1 [1 exp(-2)), have
0 s (2) E(x) (x)(2 22(x)(2 =))
for Iz x. <
Proof of Lemma A.2. See the proof of Proposition [Wu et al., 2024]. The lower bound by the convexity of e(-).
The next lemma controls the decrease of the risk L,,
Lemma A.3 (Decrease of Lạ). Suppose Assumptions 1A and IB hold. If L(wi) 1
then we have
up
L(Wr+1) L(w.) 7 (2p² + (WW Particularly, this indicates that if [(w:)) S .,क then I(We+1) S T.(w.).
-i BhL(W.))I
14

---

Proof of Lemma A.3. By Assumptions 1A and IB, we have || S and f(w;x) is 8-smooth as function of w. Therefore, for every E [nx] we have
lg:(t 1) (:(1) (/ (We+1;x;) (wr;x:))
f (w: 0 (We+1 wo); (Wz+1 - wx)
by intermediate value theorem
S || well
- p
since W4+1 W. W VVVLE since || S Lip
S p² S 1.
Then by Lemma A.2, we have
ltgit + & = - (((()))))))(((((((((((+))(++)))) (t))
+
{(()(()) + -/(()/////////////xxx W w.) + (e((()))) [W] (et1)) + =2( since (t+
V.(::.)))))) W4+1 ww) welf]((w, we41))
+ ((gg(())((()(()))))))))))))))))) ((( X1), Wi+1 - wo) + (()) (|²)) w.ll?
by Fact A.1 and the previous inequality Taking average all data points, have
112 + (2p² + L S L L.
which is equivalent
L++1 - Le S - (2pp 2 B)RLDI
We complete the proof of the right hand side inequality. The left hand side inequality be proved similarly. In detail, we can show that:
Elq. 1))= (()) (' (1)(a {()(())) (gg (. (W(;;; Wx41 Wt) 18 (e((((()) (E[/(W,,, wo+1)l-
(1)
Taking the average over all data points, we have
Tu+11 2 Lu ij(1 BTT,))|V
Now we have completed the proof.
A.2 Increase of the Parameter Norm
In this section, we demonstrate that the parameter norm, Pt. increases monotonically during the stable phase. We introduce crucial quantity, Ut, defined as the inner product of the gradient and the negative weight vector:
Vt = (VL((w), -w)).
This quantity, ltz plays a key role in controlling the increase of the parameter norm. Notably, oe appears the cross term in the expression |/ effectively characterize the increase the parameter norm.
W.W.) By managing vt, can
Recall that our loss function is f(x) == log(1 + +)) Inspired by Lyu and Li [2020], we define the following two auxiliary functions for the logistic loss:
(xx) loog((())))))))))))())))((())l
log log(1 + e e-)), E E,
((x) = (z) = log(e" - -1), I E I.
One important remark is that if we change the loss to the exponential loss, both & and will be the identity function. Since the logistic loss and the exponential loss have similar tails, our 1b(z) and (xx) are close to the identity function, i.e..
(xx ~ ((x) ~ x2 for x large enough.
Then, we have an exponential-loss-lik decomposition of T.;
L (l) = 1 - (a.(())
(7)
1=1
1=1
These two functions vb, will help us analyze the lower bound of ey. First, we list some properties of gi and here.
15

---

Lemma A.4 (Auxiliary functions of E). The following claims hold for ł, v. and t.
((r) = =--((u)
i monotonically decreasing, while to and l are monotonically increasing.
((((())) [(2)'
& ( increasing for x E (0, +00).
Proof of Lemma A.4. The first two properties are straightforward. For the third property, apply chain rule on (((x)) to get
(=)
For the fourth property, notice that
(z)===
(1 + ez) log(1 +e=
The denominator is positive and decreasing since
dr (1 - e*) log(1 +e e)] = e" log(1 + -
Combining this with the fact that positive and increasing, we have the desired result.
Besides, have the following property of This is the key lemma to handle the homogeneous error. Actually, this lemma is another way to show (xx) close to the identity function.
Lemma A.5 (Property of t). For every x E R, we have
(x) (xx)
2 + log
Proof of Lemma A.S. Recall that
((x) log(e" - 1), ((xx) e -
-
Let = x We have
(xx) log(e 1) 1)(et- -1) - -log(ew - 1)(e8 - -1
(x)
ge
c
Define s(y) = dxx
log log 2. Then,
d(x)
s(y) log((e 1)(e" 1) + +lg((y) log log
eVy
s (g) log(e" 1) el - -
evyz
Note that the sign of s' is determined by log(e" - 1). For 0 ed s s'()) S 0 and s(y) is decreasing; for e 2 2, s(y) increasing. Therefore,
min 5(y) s(log 2) 0. v)(),)))
Since log have the desired result.
Another important property of i is that can provide a lower bound for qunn (t)
Lemma A.6 bound qmin). For every? 2 O, we have
Amin (t) (( log(L.) logn).
16

---

Proof of Lemma A.6. We use (7) to get
1 - Eggmin (1)) S L 7 - ((Amms(t)) I Tt
7 1 (1mmia (t)) 2 - -log log L
Umin (1) 2 c(( (-oo(L[) logn).
by Lemma A.4
Then, complete the proof.
Now, we are ready to give a lower bound of et. The following lemma is an extension of Corollary E.6 in Lyu and [2020], where they dealt with a homogeneous model and the exponential loss; we extend this to hongeooggessus model. The key ingredient is Lemma A.5.
Lemma A.7 (A lower bound of Uz). Suppose Assumption holds. Consider U, = (VL(We), -wW) -wc). If 2nc then
v 2 b, log(2ne" L,) 2
Proof of Lemma A.7. By definition, we have
0t = (W((,)) -w:)
Σ (tw. wi)
n
Σ
e' ( 1
(Iw.f(::ii) wo) -we((we:
Σ ( - kLe
since (x(x)) s ((x) and [V]((wx))) wi) ((w,xx)) S K. by Assumption 1C
vla (1)) (V)) (9)(t)))).(t)) kL,. since e(-) exp(-x(())
Applying Lemma A.6 and Lemma A.4, we have (t()) 2 Jmin (t) 2 (( - log(nL.) = - -oog(" - -1) -log(e - 1) 2 0.
Then we can apply Lemma A.4 to get
(()((t)))) 2 ('(1 (((-log(nLo)) ((loog(nL)) = c( - log(nLe))
d' (og((n)))))
Invoking Lemma A.5, we have
log(nL))) C' log(nL)) 2 - log(nIx) + log log 2 2 - -log(nT4) + log log e = log (og(2nl..).
Putting the above two inequalities together, we have V (t)a(t)) 2 log(2nLi), for every = I....,n.
Plugging this back the bound of th, we get
(ta(()) 1og(2nL)) kL:
Vt
n
L, (2nLe) kL1
= LL log (2nc" L,) 2 0.
This completes the proof.
Right now, we get lower bound for v., which the cross term in the expression of || [|2. The next lemma controls the increase of the parameter norm Pt using 2t and L.,
17

---

Lemma A.8 (The increase of Pt). Suppose Assumptions 1A and IC hold. If L1 S
min
7(4p*+28) then
2ne
0 S 21ju: S Pi+ - 21p0r (1
2 log((nee* Le)
Proof of Lemma A.8. By definition, we have
Pe p 21 La we) + +1
2îju: + // 2 21₁u; 2 0,
where the last inequality is by Lemma A.7. Besides,
VLL 11?
Pi+1 p.
=2i
1
20,
nLL p²
2fju
1
by ll s s,, Assumption 1A, and Lemma A.7
20
L
21ju
1
by L s +44p2+20)
214
2fju
by Lemma A.7
2 (ong(2n
This completes the proof.
A.3
Convergence of the Margin
In this section, show that the normalized margin of general predictor converges in the stable phase. Recall that we define the (normalized) margin as
7(w) min4e(x) [w;;
[
However, this normalized margin is not smooth function of w. Instead, we consider a smoothed margin 7 as an eesy-to-analyzz approximator the normalized margin [Lyu and Li, 2020]
(w) == - log L(w,)
(8)
[ l
We that 7thh is good approximator of 7. We can then use 7 to analyze the convergence of the normalized margin since they share the same limit (if exists). While y is relatively easy to analyze for gradient flow [Lyu and Li, 2020], analyzing that for GD with large (but fixed) stepsize hard. To mitigate this issue, we construct another two margins that work well with large stepsize GD following the ideas of Lyu and Li [2020].
Under Assumption 1, we define an auxiliary margin as
- -log(2ng)) L(w))
(w))
(9)
|w|
and a modified margin as
(L()))
+ (4p² + 2,)
(w)
where (xx = log( - log Loge2ne"mn x))
(10)
|w|
log(2ne"*)
These two margins provide second-order correction when viewing large stepsize GD first-order approximation of gradient flow. the following discussion, we will show that 7(w) = ("()) x
(w) R (w). At last, we will use the convergence of 7f(w) to prove 7(w.) converges.
The following lemma shows that 7(w) s (w)).
Lemma A.9 (Smoothed margin). For the smooth margin (we) defined in (8) and the normalized margin (w,) defined in (3), we have
18

---

1, When L S 2n
we have
qmin (t) - -log Le S 1og(2n)) + Amin (t),
and
7(we) S (w,) S ((we) + log(2n)
Pt.
Assume Assumption IC holds. If L - O, then 170 - (we)) - 0.
Proof of Lemma A.9. To prove the first claim, notice that
1
1
L
E(gmin (1)) loog1 + exp((=mmii ammn (1))) S nhe 2
Therefore we have
Imina (t) Se² - - 1
Using _ < log(1 + x for S 1, we getet
(m)al() (t) < (((min (t)) = llog(1 + Amiaa()) $ Qmiu(t)
Then bound T by
Amina (1) _ ((Amiuult (t)mi(t))= S Le S El (t)) S e Amin (t)
2n
which is equivalent to
mmin (t) S - log Ta S log(2n) + amin (t).
Dividing both sides by Pt proves the second claim:
(ww)) == Imin (t)
= log L,
log(2n) log(2n) Amin (t)
S (we)
S (ww))
Pt
Pt
Pr
Pt
For the last claim, only need to show that Pt This is because if L, 0, have for any E [n], fff(w xx) + co. Using yi (wei:;;) I Cr will + C, from Lemma G.2, we have Pt Wx|l2 cc.
The following lemma shows that (w) z (w)).
Lemma A.10 (Modified and auxiliary margins). Suppose that Assumption I holds. For the modified margin (We) defined in (10) and the auxiliary margin (ww)) defined in (9), we have
If L 2ne there exists constant such that
7f (ww) $ (ww)) - 7(w.) < (1+
log(1/L(w))) (ww)).
Proof of Lemma A.10. To prove the first two inequalities, notice that
+ (4p² (1p(2+++ + 23)r
((L)) = lolologeee L.) exp
(Tog(2ne>
using (10)
1+ (4p² 2B) < log(2ne" Lx). since Ia S 2nex exp
< 1, and log (on(2ne" Tx) >
log (2nc L.)
- < - -og(L))) log(2n)
< gmin (1).
By argument in Lemma A.9
Using the above, (8) to (10), we have
(we) = (((((W))
log((ne) ITa) Lx) [we]]
= (/ <
Zmin
=; (ww))
|eel
wel
19

---

Then, we will prove the remaining inequality. First, we have 7(we)
7(w))
(ww,)
7<(w;)
((w)) (Wwe)
Qmin (1)
+ (4p² + 28)
By the definitions of 7, 7,79,7"
exp
(2nel L1)
logg2ne* Lt)
-log(Le)
(4p² 2
s -log(2ne* ).,)
exp
Since Ymin (1) S 1log(- Le) by Lemma A.9
(oge2 TTt)
1 + (Ap² + 2
1 log(2ne"
exp
-log(2ne* Lt)
To simplify the notation, let
1 + (4p² 2 and 02 log(2ne"). Since Le S -
log(222) (2ne* Le) 2 2 > 1, we have
1 (4p² + 23)
C1 log((ne) Txx)
S c1.
1og(2ne* Tue)
Besides, given c, we have e" S c... Therefore,
1 + (4p² + 2
C1 llog(2ne* T.x)
Ci exp(ci)
exp
log(2ne T++)
eep
11+
log Mog(2ne*1) TTt)
Plugging this into the bound for (w.)) ((we)/7((w.) we get (we)
C2 log(2ne* 1.,))
C log((nne Tnn)
+
exp
(ww,)
C2
exp(ci)c1
< + log 2ne L.)
1
log(2nee* Le)
S + C2 + exp(ci)e + 0201 exp(c1)
Since - log(onee Lt) 2
-log(2ne* Lx) =1+ C2 + exp(ci)e + C20 exp(ci)
-log L: 02
Note that - log La 2 2> 1. Because a is decreasing when x: 2 02 + 1, we have
1 llog L. Plug this inequality into the previous bound for (ww:)////ww and we get
- log Lr
02 + 1
S C2 +1
log L+ C2
log L
5(w,) - (c2 + exp(ci)e + 02c1 exp(ci) (ep(ci))(e2-1))) 1)
(w,)
llog L;
Let = (cg + exp(c)ec + 02c1 (xp(e))((2 + 1). We complete the proof.
The next lemma shows the convexity of defined in (10). The convexity will help us analyze the change of (ww,) in the gradient descent dynamics. Specifically, we are going to use the property that
((x) - 4(v)= 2 4' (z) y), for all x,
Lemma A.11 (Convexity of Φ). The function O (x) defined in (10) is convex for 0 < x < 2ne=tx
Proof of Lemma A.II. Check that
1 Iog(2ne>x) (1 + (4p² + 23)7)
V1 (x:
x loloolgnee (2ne*x)
and that
("(x)= = (1+ + (4p² + 23))) (2 + loge
(2nc" x)) log (2ne" - log (2ne
x2 log (2ne*x))
Note that when S 2 have log(2ne"
2, which implies
2 + log(2ne" x) S 0, log(2ne" x) < O, and - l (2ne"z) log(2ne" x) < 0.
Plugging these into the previous equality, have 4 (xxx) > 0 when 0 Ine-tn.
20

---

Before we dive into the proof of the monotonic increasing (w)), we show that y is bounded first. The convergence 70 direct consequence the monotonic increasing and the boundedness of 7"
Lemma A.12 (An upper bound on 75, 72, 70 and y:) When L S 2ne ((44p++8) } for 2
there exists Bo such that
log 2n
((,,) S (w,) S <w)(w,) S (w),)
+
Bo-
Proof. Apply lemma A.8, we have ||well 2 Pr 2 Ps. Then we can apply lemma G.2 and there exists a constant Co,.s such that for all i,
(u:f(W(;;;)) & |
Hence,
arg min Wi xw)
7w.))
S Cp,.k.
Besides, by Lemma A.9, we have
log 2n
log 2n
(ww)) S 7(w.) +
S Cp..x +
Ps
By Lemma A.10, we have
log
7 (w,)) S (w:)) sw:( S Cp..x
Ps
Let Bg Cp.,x + log Then, we complete the proof.
The following lemma is a variant of Proposition 5, item 1, in [Wu et al., 2024]. Before the lemma, we need some auxiliary definitions. let us define
[: Vs = 0,0,
Wt
#4 (I=-0))))
Therefore, we have
- ||
The key point of this decomposition is that we consider the gradient of the loss function as a sum of two orthogonal components. The first component v; is the component the direction w,, and the second component μu; is the component orthogonal w,. We will show that the modified margin (w:) is monotonically increasing. And the increase of >w(,)) lower bounded by a term that depends ||
Lemma A.13 (Modified margin is monotonically increasing). Suppose Assumption 1 holds. If there exists s such that
L, min
ex+22n' 7/4p2 + 2.9)
thenfor: t 2 8, we have
L1+1 < Lr.
2 2 I1 (2ne T.t) 2
p - ot 2 22j
(we4)) log (we)
pell 2 Pt+1 p.
As a consequence, (we) admits finite limit.
21

---

Proof of Lemma A.13. The first claim is by Lemma A.3 and induction. The second and the third claims are consequences of Lemmas A.7 and A.8, respectively. We now prove the last claim. By Lemma A.3, we have
L Li == L(W(+1) - L(wi)= S -ū(1 (2p² + +))LL)|\VL.1..
Multiplying both sides by 2v: 1- 2ingl
> 0, we get
+220+++88)
1 2log(2ne*L Le)
20, (Le41 Le) S 2 (1 1
/
1 (2p² + B)
2 (ogg(2eeLL))
From Lemma A.8 we have
0 2 - 12 S 27 (1
2 Mog(2ne* Tue)
Using the above we get
1 2 2og [e)
20 (Le+1 - Lo) -( - v.
(11)
1 - (2p² + B)
Recall that
| 11 =||2o|1+ +1|mee1?
For Vt, we have
| - (We, - VIn) Vt
Pt
Pt
Then we can decompose VL, |2 as
|L = || + || = + ||
(12)
Pi
2
Plugging this into (11) and dividing both two sides by 2n2, we have
1 - 2 22ne 1.,)
(
(2p² + BB7L))) (Lo+1 - L1) - - (pit1 - -p3)
12
Pi
2
20
By Lemma A.7, we have Us 2 LL (one2 L.). Define
((xx = (1- (2p² + B)
2 log 2ne
Then we have
((L,)((T411 T.,)
(1 (2p² B) Lr log (2nc" L,) (T Iee
2 Tog 2nee T.,)
1 2logl2ne" Ls)
(T4+1 Tx))
(13)
(1 (2p² +3)0I1)00
1 (0t pi+z)
(1
p
| llu.13) 12)
5
P
20g
We are going to show that Vxxx) < - 1 (x). Note that when 0 < x S min
2m 1110 }. we
have log(2ne"x)) x) 0 and - (2p² + B 3)Tx)
0. Therefore, have
1 2log(2ne"ar)
V(x) = 1 - (2p² + 9)
(14)
x log(2ne**))
=.-0
>0
To get an upper bound of (xr), we just need an upper bound of J. Let == "log[2nenz]
0 and
b = (2p² + Fp E (0, 1/2]. Then we invoke Lemma G.4 to get
= 1 < 1 2a + 26
+ (4p² + 2.
log(2ne*x)
22

---

Recall that x I ex122n I 2ne"
1 and 2ne" 2 1. Then we apply Lemma G.3 to get <
log(2nenz)
Plugging this into the bound of J, we get
J S log(2ne*xr) + (4p² 23) 5
(1 + (4p² + 23)r).
(15)
log (g((2e*x))
Plugging (15) into (14), have
log(2nemx) (1 + (4p² + 25)m)
-1
(xr J log (2
-x(xx,
log (2ne*x:)
which verifies that (x) S (x)) By this and (13), we have
1
V (T,))) L.) o (pE) pt)
+
mell
2 0,
2
20
where (xx) log x log(1 (1/r)/ Recall that for 0 2n
F(x)) is convex by Lemma A. 11.
By convexity of yp and $, we have
+ - +
log
log
+
|
2 0.
Pi+1
2
20l
By the definition of yf in (10), this can be rewritten as
1
log 1
log logn"(we+1) - >f (ws)) = (L((++1) [b(L)) +
log
Pr+1
p²
log p
- llog
|ell
pt
2v
[4 1 log
P++1
Pt
0,
where the last inequality because of Lemma A.8. We have shown that (((,)) monotonically increasing. By Lemma A.1 (w,,) is bounded. Therefore 7(w,) admits a finite limit. This completes the proof.
A.4 Sharp rates of Loss and Parameter Norm
Right now, we have already proved that (w:) is monotonically increasing and bounded, which indicates (w;)) converges. However, if we want show that 5(w)) converges, we still need to verify that T., O, which is the crucial condition for (w,,,, " (w,), (0))))"(())))) (We). and 5(w,) to share the same limit, by Lemma A.9 and Lemma A.
Fortunately, with the monotonicity of (w:),, we can prove that L, converges zero and even characterize the rate of Lt.
Lemma A.14 (Rate of L, general model). Suppose Assumption holds. If there an such that
L(w.) & min 6x+22n' 7/4p2 + 213)
then for every we have
1
S L(we) S (t- (w.)²
I(w.) + 3Gp 8)
That is, I(w:) = (1/tt) + 0 t + 00.
Proof of Lemma A.14. By Lemma A.3 and (12) in the proof of Lemma A.13, we know L, decreas- ing and
L41 L, < " | 2 <
"
1 r2
1
(16)
2 pr
23

---

We will establish an upper bound for P: first. Note that (wi) is increasing for t 2 by Lemma A.13 and (ww,) 2 (w,,) by Lemma A. By (9), we have
log(2ne L1). (w,)
S -log(2ne* LT)
5 log((ne) Le)
Pt
70(w))
(w(ws)
Combining this with Lemma A.7, we have
Dt 2
Lg log(2nc Lt) (w)))
log(2ne" LT)
Pt
(W.)
Plugging this into (16), we have
2 - L41 L, s
which implies
in(w
(2
I L4+1
2
L
Lt Li+1
Since L441 Lt
La bet1
t 12 8.
L++1 Lt
Telescoping the sum from to I, we have
1/ (w,)
1
(1
2
T
T.g
Tit
Therefore have
L S (t (w.))
Next show the lower bound the risk. By Lemma A.3 have
L++1 Le 2 77 + (TL,) 2 - VL 12
Observe that under Assumption IA.
| =
(a:(t))) Spr.
Then have
3
L141 - L. 2
1
L?
t t s.
Let L. 3ñ0² 2- L., we have L. < 3hp²
2. Furthermore, since Le decreases
monotonically, L, < L, I 1. The inequality becomes
İ L, -
1 and apply Lemma G.1, we have for any t 2 Therefore, let c L.
L, C 2(1- 8)
This is equivalent to
L, + 3rp s)
We have completed the proof.
Furthermore, can characterize the order of pt in the stable phase.

---

Lemma A.15 (Order of Pt in general model). Suppose Assumption I holds. If there is s such that
e*+ 7(4p2 + 28) }
L(w.) I min
the t we have
= (logg(t))
Proof of Lemma 15. Note that 7f(w)) increasing for / 2 by Lemma A.13 and (w 2 (w()) by Lemma A.10. Therefore,
log 1og(2me= L.) - -log(2ne* L.) - -oog(nee Lt) (((w))
(we) Pc
(w.
Combining this with Lemma 14, we have
< log /L(w,)-33)p't (log(TT)).
2ne6
7(w,))
Besides, we have Juin Therefore we have
(log L - -log by Lemma A.9 and Juin S Bupt by Lemma A.12.
allog 1.11 log nhe
log (t (t=s)i)((w.)²
Pt
2n
((log(tt),
Bo
2Bo
2Bg
where the second inequality is because for ((x) defined in (6), ((x) 2 z/z for x 2 0.6, and the third inequality is by Lemma A.14. Combining them, we get
Pr = (log(t)).
This completes the proof.
A.5
Proof of Theorem 2.2
Proof of Theorem 2.2. We prove the items one by one.
The monotonicity of L, comes from the result of Lemma A.3 directly.
Item is due to Lemma A.14.
For item 2, the monotonicity of Pt comes from the result of Lemma A.8 and the order is due to Lemma A.15.
For item 3, we first know that L, - 0 by Lemma A.14. Then, by Lemma A.13 and
Lemma A.12, we know that >((wi) converges. Combining these with Lemma A.9 and Lemma A. 10, we know that (w::) is an (1 + 0(1/(log ()) muutiplicative approxi- mation ofry (w,), and 7(w,) shares the same limit as (ww.)
B EoS Phase Analysis
In this section, we focus on the linearly separable case, that is, we work under Assumption 3. We mainly follow the idea of [Wu al., 2024] for the proof. In detail, consider comparator
u u + u2
where where
(1)
log(>2n1I) + n
with (
a
W., I,...,m m,
(17)
ar
25

---

and
12
with ub) 27 Was j= I......, m.
71
(18)
(m)
Consider the following decomposition,
||r+1 u|2 |W, - ul12 + 2mm (W()), u w.) m VL )|
|W - -ul? + 2 (L(w,), u1 - w.) +mar ( 2(V)((w)) uz) + myyl m VL(w)))l
=:/(we)
We aim to prove 1 (w,) v + - L(w,) and I2(w,) S O. Then we can get a bound for the average loss by telescope summing the decomposition. Here also introduced the following vector w.:
W. =
On W.
We can observe that u1 logln (1))))))))-- w. and U2 27 W..
Lemma B.1 (A bound on I (w) in the EoS phase). For u1 defined in (17), we have
1 2TT
4 (w) = (VL(W), in - w)
L(w).
Proof. Since L is averaged over the individual losses incurred at the data (x,, yo y= and gradient is a linear operator, suffices to prove the claim assuming there is only a single data point (x, y). Then by Assumption 3, we have
(yx, w.) 2 7 > O.
Then the loss becomes
L(w) = ((uf(w::))
= yー Σ
1 m
ajo(x w())))
Now we expand I1 (w):
(w) (VL(w), un - ww)
('I((w::))) /VV[(:: x), -w)
(x(((wx
axyo (x (()))x (u, (1) - w())
= =[([f(w:)))
a+ 6'(x T Tx + o(x T (())) 6'(x (x))))x w(a))
axyo(x T w())))
(19)
zids
By definition we have J2 yf(w; x). As for Jr, using ø' 2 and agya T we have
u (k) 2 by Assumption
0' (x ww (x))x ufth (A) + o(xx T TwTw"! (()) - 6'(x" Txx (k))) T w (A)) a
m
17
m
(k)
a+d 6(x (6xTTw(i)) w(k)) o' (x wil)) )x wilh)
0gooyX
u
171
m
26

---

1
2 1og(?2nT)) +
yx yx"w. 1 Σ ax |x
R 0
0
k=1
since 10(x' T w(h)) - 0'(x" (wx](( T w(()) 1k)
S x by Assumption 2C
11g()?'1I') + K - K
m
since yx T w. 2 7 and
=1og(?2/TT).
(20)
Plugging in J/2 yf (w; and (20) into (19), we get
(1(w) = (VL(W), i1 - w) = (yf((w: - J2) s (u] 1ogg(??/7)= v[(;;))
since e' <0
s <(log(1> - ((uf(w:x))
since E is convex
L(w).
72n
where in the last inequality, we use ((x) S exp(-x) and we only consider a single data point. This completes the proof.
Lemma B.2 (A bound on (w) in EoS). For u2 defined in (18), for every w.
12 (w) = 2(V u2) + mnn V VL( 1|² s O.
Proof. For simplicity, we define
((w()) = 9(w())==((y (y: (x] (xTw(()). Note that -1 - (() < 0 and 0 < a ()))) S 1, we have
-1 S g: g.( SO.
Under this notation, we have
@L(w)
El (w: 2x((x w())xx
Ow;
i=1
9i ( amm y.x.
So have
12(w) = 2LV:2((L((w.)) u2) + mTT V(())l
m
27
Di (w (a)aa
x3 xTTu + + Σ gil (i( )a))))))),
m
For the term inside the bracket, we have
2
gi ( x3 uu
( w. +
+
x2
gil (i(w()) )azyixi
n
2
Σ (w
(1) )ayllk
T 11 x,
a
-
1
Σ ge (w()) )a
since u2 24
(1) naj
w. by (18)
27
2
Σ (w
(1) |03
2 17
7
+
Σ gi( (1) )a
since gx(+) < 0 and y/ w 2 7
27
ト (
go
(W) +
since az = 1
27

---

S/
Σ Bi(w()))
1 +
Σ of w()))
since 19:(+) S and | S
<0.
since -1 9:(-) S
Hence, we prove that (2(w) S
Theorem B.3 (A split optimization bound). For every 1j > 0 and U1 + U2 such that
log(72ntt +
U1
with u = 0j ory
Was j = 1,..., m,
(m)
and
(1)
112
112
with m 0gj 27
We: j= l,...,m. m.
(m)
U2
we have:
|WT
1
1 T-1
1 8lg (10822((7217)) (+21)T)/a² + 8888 88n2/0² +1
||wll mnT
L
(((((t))
5
2mnT
T
n/T
k=
for all T.
Proof. By Lemma B.1 and Lemma B.2, have
|W++1
u/|² || u||² + 2m + mpmIz(wi)
||w, ull2 2mmpl1 (w:)
|w u/|2 2mmy
221/T I(wi))
Telescoping the sum, we get
|T ul
T-1
|wo ul
L(wi) S1+
2mm
2mn
By (17)and (18), we have
|wo u/12 2/ woll + |u|
s 2| woll + |n 4 12
8m 8m log( ( 21)T2
8mk2
mn²
=2| +
a²²²
72
which implies that
W
u/l2
T-1
Σ
((w(k)) <
+ 8 2 ()2 + 8k. 1/2 + tn²
||oll mn?
2nripl
T
2/
km
We complete the proof.
B.1 Proof of Theorem 3.2
Proof of Theorem 3.2. By Theorem B.3, have
1+ 8 log (72 n/) + 8x² /02 + 17²
|woll mnjT
T k=
This completes the proof.
28

---

C Phase Transition Analysis
In this section, we will analyze the phase transition. In detail, we follow the idea of [Wu al., 2024] and apply the perceptron argument [Novikoff, 1962] to locate the phase transition time, Compare to the previous EoS phase analysis, we need an extra assumption on the smoothness of the activation function, which is the Assumption 2B.
To proceed, let define the following quantities for the GD process:
" == n 1 exp (vi((w;x.)) F(w)
11
G(w)
n exp y:f(w:::))
Due to the seff-boundedness of the logistic function, we can show that G(w), L(w), F(w) are equivalent in the following sense.
Lemma C.1 (Equivalence of F)..
1. G(w) S L(w) S F(w).
2. (r)((w) S vm V/(W)] S G(w).
3. If GfW(w) S 2n then F(w) S 2G(w).
Proof. The first claim by the property of the logistic loss. For the second one,
m Σ
[V[()]]
(uif((;;,)) y: a,m _ w (1)
n
j=1
2
8l (w))) -m-
since [|ı:aşØ(x] w()))x/l S
<
j=1
1
G2 (w).
Besides, we have
Vml V (T[(w)]] )1 2 (-VT((w) w.)
since |ww]| || S mm
(w:f((::))))''0 w. wx w.
nm
1=
2 a7
since o' 2 a and yxx, w" 27
1+ exp (w:f(w;xi))
= Gww)
For the third claim, by the assumption, we have
1
1
n 1 exp (bif(w;x))) S G((w) S 2n
which implies that
vi((w;x)) 2 O. Vi [n].
Therefore.
1 n G(w)
Σ n 2= 1 + (xp()=((w;x.))
1
1 n2
2 F(w).
(xxp():(w;;x;) n
We complete the proof.
The key ingredient of the phase transition analysis is the following lemma. The main idea is to consider the gradient potential G(w) instead of the loss function Z(w) in EoS phase. And this will decrease the order of the bound of phase transition time from Ō( to Ö(n).
29

---

Lemma C.2 (A bound of lw, []) For every we have
2 + 8 loogl1" + 8k/a + 4n
W S vm
+ 2/|lwo
7
Proof of Lemma C.2. By Theorem B.3, we have
Ww 1 2mml
Wi u 2mzyl
=u(m(t)
1 + 8 log (r2n(/? 8K² 1882//a 1²²
|woll
S
72
mnyl
Besides, we know that
4m log(>2ntt22 + 4mm?
mnn²
||ll s 2| + 2||u2|12
+
a227
272
Combining them, we have
+ 24 24 + 24x²/a² 317²
|/ 2| < 2||w, u/12 2||u|| <
+ woll.
72
Hence, we can get bound for |well-
2 8 (log(>"nt)/aa + 8k/a + An 2/ woll
|well vm
7
Lemma C.3 (Gradient potential bound in the EoS phase). For every n) have
1
Σ
G((((k))
(We, W.) (Wo. w.)
/mllwwll
(Wo, w.) w.)
t 1.
manytt
marmt
k=
Additionally we huve
+ log (2²11)//a / + 8k/O + 4n
3||woll aynt"
Σ ((w(t))
t 821 1.
2n
K=
This
Proof. This is from the perceptron argument [Novikoff, 1962]. Specifically,
(We+1, W.) = (We, w.) mm/(VI(w.)) w.) of (.(l,..((w:::)))))okx w w(()) (xe, wa)
= (We, w.) -
(We, W.)
of (v:(((();)))))
(Wes W.) + ma.mG((..))
Telescoping the sum, we have
(We, w.) - (Wo, W.)
Σ ((w(
marnt
k=0
m well (Wo, W.)
marmt
2 - 8 &log (x2 + 8k/e + 4n 3||woll marryt
+
by Lemma C.2
2nt
We have completed the proof.
30

---

Besides, we can make use of the equivalence between G and L to get a bound for the loss function which is independent of the initial margin at s.
Lemma C.4 (A risk bound in the stable phase for Two-layer NN), Suppose that there exists a time such that
L(w.) S min
1744 1 23)
2e
Then for every 2 8 + 1, we have
L(ws) S (t s)a?n?"
Proof. By Lemma A.3 and f(x is ym
Lipschitz and smooth, we have
Lb+1 I LA mny(1 - (2 + ()n((w)))
By Lemma C.I and L, S "(+28)) we have
a2n² L2
LK+1 S La
Multiplying 1 both sides, we have zo
a2
L Lx+11
<
2
1
Lx+t1
Lx
Taking summation for S,..., t
1, we have
1
1
(t
3)
2
L,
L
Lt
La
2
(t st=
At last, we will use the bound for the gradient potential to get an upper bound for the phase transition time.
C.1 Proof of Theorem 4.1
Proof of Theorem 4.1. Applying Lemma C.3, we have
2 log (>²rr)/a + 8n/a + 4n
3 woll Vmaamp
G(((())
1721pT
K=0
2 8k/a 8log(1"))/aa + (4+ 8/a)n
3 3olloll
since log(r) 1
277
Vmaym
Let Ci x22
02 = (8 + 43). Note that we have
+ 8k/a $
1
if 2 4(2 8k) C27) + cin
2721T
4(cin + c2T)
na²
8 8log(?2T)/a
<
1
if 2 2 128
c2n + an
log
C2n) + cin
since Lemma G.5
4(cin + 021)
na²
(4 + 8/a))
1
48
if 2f >
(c2) + cn))
17²
4(cin + c21)
2
3/ Woll
1
12 (c2) + an)
|oll
4(cnn + on)) if
maryje
"
Vm
and that the two conditions are satisfied because
2 128(1 4k)
max C27), Cin,e, C2n + cin log
onn + ein
(o2n + cnn) [wol
a2
1y m m
31

---

C21) ein
227 + cin
22 + cin
48
12 (c27 cin)
|wo vm
max 4 + 816)
128
log
a (cen qn).
n7a
nja
a
So there exits S such that
1
G(ws) min
n18 + 4,)
Then we have L(w.) & F(w.) S 2G(w.) S
{ cutton ((1+20)
We complete the proof.
C.2 Proof of Corollary 4.2
Proof of Corollary 4.2. The main idea is to show that S T. Note that by Theorem 4.1, we have
T 128(1 + 4k) max C27), C,n, C27) cin
C2ll + cin
(c2n tin)
|wol
log
v m
in which expression c = * and C2 = (8 + 43). We can verify that,
128(1 + 4k)
128(1 + 4x)
2222
C27
T
a²
a²
256( 1 4x)02
2
128(1 + 4x)cn
a²
22,27
Furthermore, we have n S
Hence,
356(1e
a²-TT 256(1+48)
c2n cIn
256 4
4kj
< 202
222T
22PT
256(1+1x)cz
11+
We get that:
128(1 + 4x)
Cn7 + cin
C2n cin
128(1 + 4R:)
128(1 + 4R)
log
2
C2 1n(2222 5
402 <
a²
a²
a²
2
128(1 + 4k)
(cen + an)
||oll
128(1 + 4x)
Woll
T
202
a²
7
m
a²
Vm
2
Hence, we have T S T
2. Applying Theorem 4.1, we have
2048(1 + 4x)c2
L(wr)= < a2>2n(T'- -T S
= O(1/T²)
2222nT
a47472
We have completed the proof.
C.3 Proof of Theorem 4.3
Proof of Theorem 4.3, The main idea is to construct an upper bound of 7 and apply the analysis in Theorem 2.2. Note that give Wo 0, we have
S(wo: x;) =
ax (k)) sac((0)
m
where = -m m.m
m. Therefore,
[VI(wo)](()
ak ((())x + (sa((()) 2 (sa))()) m
2
1
Qk o/ (0)x2
2
m
ak ((sab(()) (0) X1 +x2
m
as (()ss(())))0(0
v 7-
4
32

---

---1² we have
Let = ()))))(0))))(()
(k) 0 [][L[Wo)) -nax.
Therefore,
1
f(w):
axol (naxx))
m
We notice that -x 0 and
-x x
O, when y 0.1. Furthermore, we have
f (wi;
x(-x (17²)
+
1
Σ
o(x (1p²)
2x=1
(0(0) - x (
+
(1x)0'0(c)) (c2))
a+=1
m==-1
(a0(0) 7 x
1
(G) +
Σ ((2))
=1=1
a4=-1
S (a) n/xi xa
('(e)) 2a. a
Note that
} (() jxx xa)= S / [((([[])))) S L(w1) S L(wo) ((s)(())).
1
We apply Lemma G.7 to get
Isa [sa)(0) I
x xa
use C3
|+0(01+133
Now we know 7/ S C3. Furthermore, notice that
x, Rc
[VL(w:]] & I S Lo.
We get that || - well S noo S cgLo. Hence,
|/ [((w+11;;) f(w:; S cz czLo- and T 2 bo bg $ lo-
Assume that 4. min en (88488)
We know that
4. min ex+24n' ca (8 + 4,8)
la-
We want to show that there is an lower bound for L,, Now that
L 2 [[(((ww::x)) + +(f)))((((w)))
Applying Lemma G.6, we can get that
La (p)-caLLoLL--1 eexp=eexppppcc exp(-caLo)L,
Recall that by Lemma A.14, we have
L
t 2 8.
I, + (1
Combine this with p Vm 17 mm and L, 2 exp(-c3Lo) U4B and we get
L explcaLo) + 3t - s)*
1. 12 8.
Note that when 5 L, 2 4-- We can get lower bound for Tae by
C4
L 2 exp(es l
+ 3rt
exp(csbo) + 30st 1
explosLo) + 30st
t
L.
where C4 exptog Mol +30g
depends only on {a, 1 (00), K, and n.
16
33

---

D Scaling and Homogenous Error
In this section, we consider different scaling of two-layer networks. We add scaling factor into the model, i.c.
f(w: x) = Σ o ==1
(0 w())).
We will show that given limited computation budget 7 (total iterations), larger b and a corresponding best rj 7) m will achieve the same best rate as 1, i.e., O(1/T²). While for smaller b the rate is Ob O(b-3/(T")) )) Before we present the analysis, here are the bounds with b and ij the process of Lemma C.3:
following
Σ I(wx) 1 + (7")((( + 8k2/n²² 1 +
|wol mnt
Pnp
G((+) 2+ +++8lg((22n)//11 (ab) 8k// 2nb
3 3
ar
/mnbt
Case when 2 1. Given the previous bounds, we have the following results following the idea in Appendix C:
Gradient potential bound: G(w,) S C for all & 2 O.
:
Phase transition threshold: G(w.) min ( (1/e 1/n0n/1/n(1p22 + 4.b)].
Stable phase bound: L(w:) S CbAn(t-s)*
where C' depends a, Combine the first two arguments and assume </8ppb2 4,36) 2 4 We get S Cr)(pp222 + 4,86). Plug this into the third bound. We have
2
I(WT) S Cb2 Cn(8p2b2 + 4ßb))
It's obvious that the best 16+25CC+8300*
T Hence.
8(8p262C' 0
L(WT)= <
C12T
T2
Then, the rate is still O(1/T²).
Case when b 1.
Similarly, get the following bounds:
Gradient potential bound: G(w.) < Cb-2 for all 2 0,
Phase transition threshold: G(w.) S min 114c 1/n(8p212 + 436)
Stable phase bound: L(w.) (b-
where C' depends on a. 7. Without loss of generality, we can assume 1/pppb + 4,36) 2 4 since 1) can be small enough. Then, we have
Cn(Spp + 43b
Then, we can get
L(Wr) < Cb2 Cn(Spp 4/5-11)
It's obvious that the best 16p Hence,
8(8p2 436 134--3300 3C)
b-a
L(WT) <
CT
7
34

---

Combining the analysis for two cases, we observe that when b 2 1, the fast loss rate is O given finite budget T. While b < the rate O ((b-3/(6) In our main results, we b for the mean-field scaling. Under the mean-field regime, all bounds are independent of the number of neurons since we consider the dynamics of the distributions of neurons. Alternatively, if we set. b = /m, then the model becomes:
f(w; x)
Σ ajo(x w(1).
vm
The model falls into the NTK regime. The loss threshold will be related to m. but the loss rate is the same as that of the mean-field scaling.
E
Extra Experiments
Here provide additional experiments to support theoretical results. In Figure 3, show the test accuracy two-layer networks for CIFAR- 10 under the same setting of Figure 2. We can observe that large stepsizes lead to stronger implicit biases with "nicer" features.
In Figure 4, we show the training loss and margins of two-layer network with leaky softplus activations on a synthetic linear separable dataset. We can observe that both neurons have negative margins during the training, while the network's margin increases and becomes positive. This indicates that even the two-layer networks can have complicated dynamics. It remains open question to understand each neuron's 's dynamics deep networks.
0.80
0.75
Test accuracy
0.70
0.65
ñ
0.60
ñ =
ñ = 0.5
0.55
ñ = 0.25
2000
4000
6000
8000
10000
Figure 3: Test accuracy of two-layer networks for CIFAR-10 10 under the same setting of Figure 2(d)-(f).Th results support our intuition that large stepsizes lead stronger implicit biases with "nicer" features.
Additional Proofs
F.1
Proof of Example 2.1
Proof of Example 2.1. Recall that the two-layer neural network is defined as:
f(w;
Σ axx wi))).
We can verify that if ((z) is 8-smooth and p-Lipschitz with respect to then f(w;x) is 8/m-smooth and p/ Vm-Lipschitz with respect to w. This because:
(((w)= Be' (yf (w:( f(w:
35

---

Loss
Normalized margins
10a
10F
101
((W)
L(Wa)n
10"
-0.6
Neuron 1
Neuron 2
10
-0.8
Model
Maximal margin
10.2
-1.0
10000
20000
J0000
40000
50000
10000
20000
30000
40000
50000
Iterations
Iterations
(a) Training loss, synthetic dataset.
(b) Normalized Margins, synthetic dataset.
Figure 4: Training loss and margins of a two-layer network with leaky softplus activa- tions on a synthetic linear separable dataset. There are five samples in the dataset, which are ((0.05, 1, 2). 1). ((0.05, -2, 1). 1), ((-1, 0, -) -1), ((0.05, -2 -2), 1). ((0.05, -2). 1). The max margin direction is (1, 0, 0) with normalized margin of 0.05. The network only has two neurons with fixed weights 1/2 and -1/2. The leaky softplus activation is δ(x) (ær + <(x))/2, where is the softplus activation. The stepsize is 3. We can observe that both neurons have negative margins during the training, while the network's margin increases and becomes positive.
V2L(w)= = (f"" 02
and that
0 (x w (x)))
V[(w;x)= = 2x (x w(())x
V²I(w;))= =
0
Now, we will focus on the parameters of each activation function.
GELU. ((x) = x erf(1 + (x/v2)) x F(x).
d'(x) F(x) + f(x),
o"(x) 2z(x) + x f((),
where F(x),f(z) are the CDF and PDF of standard normal distribution. Note that ff(x) = V 0- 22 and (xf(x)) = -= - -=)e===²// We can find the maximum of ff(x)) is V2n
(1=
1/2, Besides, we know that F(r), < and x f'(x) S 0. Combining them, we 1/2 2 2n and 3 = 2. For K, Ø d'(x)æ = =x f(z). So the bound of
have p = 1
is e 1/2/V2TT.
Softplus. d(x) log( 1 + e*). Therefore,
ex
1 S 1, x'(x)
e"
(x)x)
(1 + c00s S 1.
Besides,
0²1 )
cxx
(6(z)- - 6'(x)x) [lg(1+
1 ez
(1 ==)²
So the maximum is ((0) - o' d'(0)0 = log Besides, when x > 1, [(x))
2 r >
ez When 1t
I -0, Φ(x) ('(x)z + 0. Therefore, N log 2.
36

---

Sigmoid. ((x) = 1/(1 e) Hence,
6=
((())
(1+0-2) 51,
e2x c=
0'(x)
(1 (1+e S1.
As for K, we know that
16(x) - (())) lte x x-
1 1te lafe
(1+--0)2
(1 + 6622
Note that |arle- S 56- + 1. We have
10(x) S 1+0 +e 22 +1
(1+ 2
Tanh, o(x) 6--e-"
t= S 1. Note that
d'(x) =z) -1=)== <
('(r) = 20(x)² ((x) 2
Besides, we know that
4|x|
[ (E* + c 2)² S4.
Hence
10(x) - d'(x)l S |o(x) + ( (x) 55.
SiLU. Note that
6'xx) 1 1te + re
+ E-
0' (2 - =)
xC 22
+
x)²
(1 e e=
Because |0le S e 2z + and |a|e 2x S e 3x + We get (d'(x) s 2 and x"(x)] s 4. At last,
|z|e- (1 + e- $1 1.
16(x) x (+)
Huberized ReLU. It's obvious that d'(x) < and 1/h. Note that & is second-order differentiable. At last,
0, S x S h, x2/2hh
16(x) - n6 xxb)) x0'()))
11/2
x > h.
Hence, it's upper bounded by h/2.
F22 Proof of Example 3.1
Proof of Example 3.1. Because for activation functions in Example 2.1, 3 < 4 and S 2. Hence, for )(x) = cr + (1 - c)4(x)/4, 3 1 and 1. Besides, since 0.5 < c < 1, we must have
(ō(x))' > 0.25.
37

---

G
Additional Lemmas
Lemma G.1. 1f 2 I 2 1 and Ln 1 T, we have
L2 2 C 2
Proof. For function g(x) - (x) 1 2z. If S 2.2
then g(x) is increasing. Then
c-1 c² + c 1
s(T11) 2 gو c² e 2)
c+2
Lemma G.2. Given continuous function f(z) s.t. 1f(z) - (Vf(z), S K, then for a fixed
constant r > there exists Cr.s and C s.1. for any |||||| 2 r
[f(x)] - Cra
and for any
[f(x) S Cr,el Cr.
Proof. Since f is continuous, let
C max fxx)///
||||=
rz and consider g(s) [lay) Then we have Now for any || > r, let y Tall
g((s) (Vf(sy), sy) f(sy)
s2
Therefore, - S gs() S 5. Let = ||=11/rr
[xxr x
g(s) =g(1) + [ g'(t)dt
s g(1 + S 1z g(L) +
SrC, + k.
Therefore, f(x) I (C, F). || Similarly, show that -f(x) (C, ||x]]] Therefore, for any ||>||| 2 r,
If(x)) S (C, + - --|||-
Let D max a| St l/(x), have for any s.
If(x)) < SC, ) || + D.
We have completed the proof.
Lemma G.3. Fixing e > L, then for every 0 < < 1 we have
log cz)
Proof of Lemma G.3. This is equivalent to show that
log(cr) 2 -1
Let s(x) log(cz), then s'(x) = = log(cr). Hence s(x) is decreasing when 0 and is increasing when The minimum of s(x) is achieved at 1. which
s(1((ce) = 2 -1
This completes the proof.
38

---

Lemma G.4. Given 0 < <b
1 and > O, we have
2
1 a
1 S (1 + 2a + 26).
Proof. This is equivalent to show that
(1 a1+ - (1 2a + 21) - bb = +1+ 2a + 2ab 2b²
This is equivalent to
2b(a + b) S (a + b).
Since a + > 0 and < 2. this is true.
Lemma G.5. Given > 6, we have for any x 2c log
log 1
Proof. It's equivalent to show that clogg 2 0. Let g(x) clog g'(x) 1 c/x. When x > 2clog c > 2c, g (x) < 0. Hence, the minimal is g(2clog c). Note that
g(cclog c) = 2c log - clog e clog - log log clog - e log c clog loge c = clog c
log
2/y > 0 when y > e. h()) = 2-22 > O. Hence h(c) > h(e) > 0 and g(2c c) > 0. This leads to g(x) > 0. Then, we complete the proof.
Now we want to show that e > 2 log c. Let h(y) = log y. h' (1)
Lemma G.6. Given f(x) log(1 e= and 0, we have for any x.
f(x + c) 2 exp(-c)((x).
Proof. Let g(x) = f(x + exp( (xp(=c)((xx. Then, we have
-1
1
g'(x) + exp(x + c) exp(c) + exp(x +
0.
Therefore, g(x) is monotonically decreasing. When co, have
lim ((x) lim [f(x + c) exp(-c)((x=) exp(-x c) exp(-c)exp((sp 0.
Therefore, g > for any Now, complete the proof.
Lemma G.7. Assume {(x) log(1 + t If [f +c) I 20(x), we have
< In + |r].
Proof. Note that
1 c
(c log
1 2 + 22x S
Then,
1 + a0o
51 = c552 c0 52 clml S 33e
1 + 26 + ela
Therefore, S 1 + (z).
39

---

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

---

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

---

reproducible, they Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: Our experiments are easy to reproduce. and we provide all the necessary information in the paper. We focus on the theoretical analysis in this paper, and the code is not provided. The dataset CIFAR-IU 10 can be found online.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
Seelliicguuddee/ooeSSubbaee wtt.ssbbmssiiooooooooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: gwwi//////////eg// /aaeeee///oeeSmaaaacaa
/ //wps.cc...
onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permirted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All the training and test details are provided in section 5.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Our experiments are completely deterministic and it's not necessary to plot er- ror bars. We provide the convergence plots of different models and show that the convergence rates are consistent with our theoretical results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
42

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτor rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification: We focus on the theoretical analysis in this paper, and we only provide some illustrative experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https //tee///se.../eeeeeee cc cctppblic//eeeeecccccecc ctwt:ettthhiiieeeeeeeeeeeeeeeeee
Answer: [Yes]
Justification: We have reviewed the NeurlPS Code of Ethics and ensured that our research conforms to it.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., - if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This is theoretical paper, and do not discuss societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
43

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release any data or models that have a high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We only use the CIFAR-10 dataset. It does not have any license restrictions.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
44

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We don't introduce any new assets in this paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
45

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
46


Reviews:
[{"review_id": "jLyauHXxi3", "reviewer": "Reviewer_dNN4", "summary": "This work studies the phase transition (from EoS phase to stable phase) of GD with large step sizes for training two-layer networks under logistic loss. Specifically, the authors proved the following:\n\n- If the empirical risk is below a threshold depending on the step size, GD enters a stable phase where the loss monotonically decreases and the normalized margin nearly monotonically increases.\n- For linearly separable datasets, GD with an arbitrarily large step size exits the EoS phase due to the convergence of the average loss across iterations. Moreover, a tighter bound on the phase transition time is also provided.\n- For linearly separable datasets, GD with an appropriately chosen step size achieves an accelerated convergence rate of $O(1/T^2)$.", "questions": "- Equation (5) in Theorem 2.2 seems to implicitly make assumptions about the data structure/distribution and the number of neurons. Can you provide any non-trivial examples other than linearly separable data? I’m wondering if this theory covers XOR or k-parity datasets, as these could be benchmarks to see the separation from the kernel regime. For instance, see the following papers:\n\n[M. Telgarsky] Feature selection and low test error in shallow low-rotation ReLU networks. ICLR, 2023.\n\n[T. Suzuki, D. Wu, K. Oko, and A. Nitanda] Feature learning via mean-field Langevin dynamics: classifying sparse parities and beyond. NeurIPS, 2023.\n- In Figure 1(c), GD with a small step size of 0.02 seems to achieve the best test accuracy at the beginning phase. What happened?\n- Can Assumption 1-C be relaxed to $\\kappa \\leq 1$?", "limitations": "The limitations have been well addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This work makes significant contributions by extending existing results. The authors proved convergence with an accelerated convergence rate in the stable phase, whereas [Wu et al. (2024)] treated the linear predictor. Additionally, the authors proved margin improvement with non-homogeneous activation functions, while previous work focused on small step sizes and homogeneous activation.", "weaknesses": "- Despite studying two-layer neural networks, the main theorems require linear separability of datasets, except for Theorem 2.2.\n- The paper missed some relevant works. For instance, [D. Barrett and B. Dherin] showed that gradient descent (in discrete time) optimizes the loss plus gradient norm and studied the modified ODE capturing this property. A stochastic variant of this dynamics was also studied by [Q. Li, C. Tai, and W. E]. Additionally, [M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] and [Y. Ren, C. Ma, and L. Ying] studied the benefits of large learning rates as well. Regarding optimization in the mean-field regime, [F. Chen, Z. Ren, and S. Wang] and [T. Suzuki, D. Wu, and A. Nitanda] proved the convergence of mean-field Langevin dynamics in the finite-neuron setting.\n\n[D. Barrett and B. Dherin] IMPLICIT GRADIENT REGULARIZATION. ICLR, 2021\n\n[Q. Li, C. Tai, and W. E] Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations. JMLR, 2019.\n\n[M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] SGD with Large Step Sizes Learns Sparse Features. ICML, 2023.\n\n[Y. Ren, C. Ma, and L. Ying] Understanding the Generalization Benefits of Late Learning Rate Decay. AISTATS, 2024.\n\n[F. Chen, Z. Ren, and S. Wang] Uniform-in-time propagation of chaos for mean field Langevin dynamic. 2022.\n\n[T. Suzuki, D. Wu, and A. Nitanda] Convergence of mean-field Langevin dynamics: time-space discretization, stochastic gradient, and variance reduction. NeurIPS, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your support. \n\n---\n\n**Q1**. Despite studying two-layer neural networks, the main theorems require linear separability of datasets, except for Theorem 2.2.\n\n**A1**. We acknowledge that linear separability is a strong assumption. We use this mainly as a sufficient condition for two-layer neural networks, regardless of width, to reach the initial bound of the stable phase when employing large stepsizes. Additionally, our stable phase results do not need this assumption. We believe the linear separability condition can be relaxed, which is left for future work. We will comment on this more in the revision. \n\n---\n\n**Q2**. The paper missed some relevant works. For instance, [D. Barrett and B. Dherin] showed that gradient descent (in discrete time) optimizes the loss plus gradient norm and studied the modified ODE capturing this property. A stochastic variant of this dynamics was also studied by [Q. Li, C. Tai, and W. E]. Additionally, [M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] and [Y. Ren, C. Ma, and L. Ying] studied the benefits of large learning rates as well. Regarding optimization in the mean-field regime, [F. Chen, Z. Ren, and S. Wang] and [T. Suzuki, D. Wu, and A. Nitanda] proved the convergence of mean-field Langevin dynamics in the finite-neuron setting.\n\n**A2**. Thank you for bringing our attention to these works. We will cite and discuss them in detail. Specifically, the stepsize considered by [D. Barrett and B. Dherin] and [Q. Li, C. Tai, and W. E] are small or even infinitesimal, whereas our stepsize is large and causes EoS. The work by [Y. Ren, C. Ma, and L. Ying] studied special linear networks under MSE loss, while we focus on two-layer networks under logistic loss. The work by [M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] studied the effect of large stepsize through experiments while we take theoretical approaches. The works by [F. Chen, Z. Ren, and S. Wang] and [T. Suzuki, D. Wu, and A. Nitanda] demonstrated the convergence of mean-field Langevin dynamics for two-layer networks with a finite but large number of neurons. In comparison, although we adopt the parameter scaling from mean-field theory, we use a different proof technique that allows any number of neurons. \n\n---\n\n**Q3**. Equation (5) in Theorem 2.2 seems to implicitly make assumptions about the data structure/distribution and the number of neurons. Can you provide any non-trivial examples other than linearly separable data? I’m wondering if this theory covers XOR or k-parity datasets, as these could be benchmarks to see the separation from the kernel regime. \n\n**A3**. This is a good question. In fact, we only need to ensure that the dataset is realizable, that is, there exists a parameter direction $W$ such that $\\inf y_i f(a W; x_i) \\to \\infty$ as $a\\to \\infty$. This indicates that the loss can be made arbitrarily small. Since XOR or k-parity datasets can be realized by two-layer networks (with Softplus, for instance), Equation (5) can be satisfied. We will discuss this in the revision. \n\n---\n\n**Q4**. In Figure 1(c), GD with a small stepsize of 0.02 seems to achieve the best test accuracy at the beginning phase. What happened?\n\n**A4**. You are correct that GD with a small stepsize achieves the best accuracy early (after about 20000 iterations). We think this is because of a mild overfitting since we only use a small dataset of 1,979 samples, while the MLP has about 30,000 parameters. \n\n---\n\n**Q5**. Can assumption 1-C be relaxed to $\\kappa\\le 1$? \n\n**A5**. In Assumption 1-C, $\\kappa$ can be any fixed nonnegative number. \n\n---", "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": "ip3l8mLrP9", "reviewer": "Reviewer_RDnn", "summary": "This paper studies GD for nearly-1-homogeneous neural networks with large stepsizes. It provides two main results.\n \n\nThe first one describes the late, *stable phase* of training and can be seen as an extension of Lyu and Li (2020) result to the large stepsize, nearly homogeneous setting. Yet it comes with weaker conclusions, showing that if the training loss becomes smaller than some threshold at some point, then afterwards:\n     - the training loss converges to 0 at a $\\frac{1}{t}$ rate\n     - the normalized margin \"increases\", and hence converges\n\nThe second one studies the early stage, *Edge of Stability* phase for linearly separable data for a more restricted one hidden layer neural network architecture (with fixed output weights) and shows that during this early stage, the loss decreases at some rate **in average**. This then allows to provide a result illustrating the advantage of large stepsize (with an optimization error of order $\\frac{1}{T^2}$) vs small stepsize (optim error $\\frac{1}{T}$).", "questions": "- Is there a specific mathematical challenge (eg wrt Lyu and Li) to prove a KKT point convergence type of result?\n- Same question for extending $1$-homogeneity to $L$-homogeneity for any positive integer $L$\n- Can we say more about the normalized margin in the restricted setting of Theorem 4.1 ?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "The study of GD with large stepsizes is of high relevance. This paper puts in perspective the phenomenon of edge of stability and illustrates on a simple example how it can be beneficial to faster rates for the training loss. The claimed results seem sound and the paper is nicely written.", "weaknesses": "The provided results and settings might be a bit too weak in my opinion. As a major drawback, the authors insist a lot on the importance of large step sizes for implicit bias in the paper. However, the only implicit bias result is claiming the normalized margin is increasing in the late stable phase. As commented by the authors in the paper, Lyu and Li provided a KKT point convergence result in their work, and I was hoping the same conclusion could be possible here. As a consequence, this paper mostly studies the convergence rates of the training loss, which I find less interesting from a personal point of view. \nActually, discrepancies in the normalized margins can be seen in the different experiments of the paper, but I think it should be discussed more. In the light of the literature on EoS, large stepsizes do not only lead to faster convergence rates, but also to a stronger implicit bias towards \"nice\" (e.g. sparse) features. Figure 1b) does not support this claim here, but 1c) still seems to suggest that larger stepsizes can get better test loss. In consequence, I would have also liked a similar comparison of the test accuracy in figure 2 for CIFAR-10. \n\nAdditionally, the setting focuses on nearly-1-homogeneous parameterization. While the near homogeneity assumption is mild and really nice, considering 1-homogeneous networks is very restrictive and simplistic in my opinion. Again, Lyu and Li provide a general result for $L$-homogeneous parametrizations. Having a $2$-homogeneous assumption (or even general one) would largely improve the current work. \nActually, combining this assumption with the derivative condition (Assumption 2.A) makes the considered parametrization nearly linear. As a consequence, I fear that the proof of the second main result, which depends on these assumptions, does not significantly differ from the linear regression case, from a high level, abstract point of view.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. We address your comments below. \n\n---\n\n**Q1**. Is there a specific mathematical challenge (eg wrt Lyu and Li) to prove a KKT point convergence type of result? \n\n**A1**. Good question. The key difficulty is that KKT points with respect to an optimization problem are not even well defined for non-homogenous predictors, so extending the KKT point convergence analysis Lyu and Li is challenging. \n\n\nSince our nearly homogeneous function $f(w;x)$ is asymptotically homogeneous, we conjecture that, under suitable conditions, GD may converge to KKT points given by a proxy homogeneous function \n$$\n\\tilde f(w;x) := \\left [ \\lim_{t \\to \\infty} \\frac{f(t w;x)}{t} \\right] \\cdot ||w||. \n$$\nThere are still two obstacles to proving the above conjecture. \n\n1. Verifying the (sub)differentiability of $\\tilde f(w;x)$ is nontrivial. \n2. The GD dynamic for learning $f$ is different to that for learning $\\tilde f$. Specifically, there are examples such that $|f(w)-\\tilde f(w)|$ is uniformly lower bounded by a fixed constant. Then the GD dynamic for $f$ is away from that for $\\tilde f$, and might not satisfy the corresponding $(\\epsilon, \\delta)$ KKT conditions in [Lyu and Li]. \n\nWe believe additional regularity assumptions about $f$ are needed to show KKT convergence. Nonetheless, we think this is a great question, and we will discuss it in the revision.  \n\n---\n\n**Q2**. Is there a specific mathematical challenge (eg wrt Lyu and Li) to extending 1-homogeneity to \nL-homogeneity for any positive integer? \n\n**A2**. As mentioned in Line 300, this is an avenue left for future exploration. Extending our results from near 1-homogeneity to near L-homogeneity will require non-trivial modifications to both the margin functions and the stable phase conditions outlined in Theorem 2.2. Furthermore, verifying that the proposed stable phase conditions can be satisfied (even in special cases) could be challenging. We will discuss this in the revision. We think our margin improvement results in this work, which extends prior results from 1-homogeneity to near 1-homogeneity, is already a significant step toward addressing the general case. \n\n---\n\n**Q3**. Can we say more about the normalized margin in the restricted setting of Theorem 4.1?\n\n**A3**. Good question. Although we consider linearly separable data, the predictor is nonlinear and, therefore, can achieve a normalized margin larger than the maximum $\\ell_2$-margin. Additionally, we observe that in practice, while the normalized margin tends to be positive and increasing, the normalized margins for individual neurons can stay negative (please check Figure 2 in the pdf). A full characterization of the normalized margin is technically challenging even for linearly separable data. We will comment on this as a future direction in the revision. \n\n---\n\n**Q4**. The provided results and settings might be a bit too weak in my opinion… the only implicit bias result is claiming the normalized margin is increasing in the late stable phase…. In the light of the literature on EoS, large stepsizes do not only lead to faster convergence rates, but also to a stronger implicit bias towards \"nice\" (e.g. sparse) features. Figure 1b) does not support this claim here, but 1c) still seems to suggest that larger stepsizes can get better test loss. I would have also liked a similar comparison of the test accuracy in figure 2 for CIFAR-10. \n\n**A4**. We respectfully disagree that our results and settings are weak. In fact, extending the implicit bias results for homogeneous predictors to general non-homogeneous predictors is an open problem listed in [Ji and Telgarsky, 2020]. Our work, by establishing the margin improvement results to general near 1-homogeneous predictors, partially solves this open problem. Furthermore, we also extend prior neural network analysis for GD with small or infinitesimal stepsizes to an arbitrary large stepsize, which is more relevant to practice. We believe the contributions of this work are already significant. \n\nWhile understanding the generalization benefits of large stepsizes is very interesting, this question is beyond the scope of the current paper, which focuses on margin improvement and fast optimization.  \n\nPlease see the attached pdf for  Figure 1 reporting the test accuracy for CIFAR-10. We will include it in the revision. \n\n---\n\n**Q5**. ….Actually, combining this assumption with the derivative condition (Assumption 2.A) makes the considered parametrization nearly linear….\n\n**A5**. We would like to point out that Assumption 2.A is a sufficient technical condition to enable the EoS phase analysis. However, we only need Assumption 1.A for Theorem 2.2, which does not assume a lower bound on the derivative and allows the predictor to be highly non-linear.\n\n---", "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": "l0Wh842RPK", "reviewer": "Reviewer_PHJ6", "summary": "This work analyzes dynamics of large step-size for GD under logistic loss for non-homogenous two layer networks. They characterize two phases, first in which empirical risk oscillates and then monotonically decreases in the second phase. Additionally, they show\n1) normalized margin grows nearly monotonically in the second phase\n2) in the oscillatory EOS phase the average empirical risk decreases under certain conditions\n3) With a larger step-size, the GD undergoing phase transition is more efficient.", "questions": "1) Will the loss landscape with non-homogenous layers be the same as the linear networks case? For example, the authors show that in logistic regression, the landscape is a valley instead of a quadratic basin. The iterates converges quickly in max-margin subspace direction and the oscillations within the orthogonal direction becomes small progressively with time. Can you give an intuition of how the loss landscape may look for non-homogenous layers? \n2) Is it true that unllike MSE losses, where there is a chance of divergence for very high lrs, the logistic regression losses always converge irrespective of how large lr is used.\n3) I missed the part in the paper where the authors discuss progressive sharpening. Usually for MSE losses, PS first takes place unless the sharpness hits 2/lr after which loss oscillates around walls of basin. Does logistic regression loss do not exhibit PS? Seems like EOS starts from the offset and then oscillations stop when iterates enter a valley of flatter oscillation directions.", "limitations": "As discussed the analysis breaks when considering non-differentiable activation functions like RELU. The analysis presented in the paper is an extension of previous work which I mentioned with some relaxation on the non-linearity (see the three conditions). It would be better if the authors can emphasize this limitation and point out the novelty directions in this paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) The paper is well written and builds on top of exisitng works on EOS for logistic loss https://arxiv.org/pdf/2305.11788 and https://arxiv.org/pdf/2402.15926 and extends these two works from linear network to two layer non-linear network. \n2) This paper analyzes the large learning rate regime with non-homoegenous two layer networks (for logistic loss). This is a practical setting since most works study gradient flow or linear overparameterized models.", "weaknesses": "1) I liked the way the authors introduced the three model conditions to extend their result from Theorem-1 in https://arxiv.org/pdf/2402.15926 in terms of the three constants (Lipschitzness, smoothness and near-homogenity). However, these conditions do not hold for Relu network (assumption 1.B and 1.C breaks), which probably is most used non-linear activation out there. Infact there are a bunch of non-differentiable activation functions for which these assumptions do not hold as well as the final result. \n2) By comparing the proofs of this work with that of https://arxiv.org/pdf/2402.15926 , it seems like similar proof technique can be utilized to incorporate the three assumptions to derive the same results involving the three constants.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. We address your concerns below. \n\n---\n\n**Q1**. However, these conditions do not hold for Relu network (assumption 1.B and 1.C breaks), which probably is most used non-linear activation out there. Infact there are a bunch of non-differentiable activation functions for which these assumptions do not hold as well as the final result.\n\n**A1**. Note that Assumption 1.C holds for ReLU networks with $\\kappa=0$ when replacing the gradients with subgradients. \n\nSince our focus is non-homogeneous networks, we choose to work with differentiable predictors to simplify the analysis. While we do not cover ReLU activation, we have addressed many other commonly used activation functions, such as GeLU and SiLU. These two are especially interesting due to their non-homogeneity. Therefore, they cannot be covered by the prior theory. Given that our work is the first to prove margin improvement for non-homogenous predictors, we believe our contributions are very significant. \n\nExtending our results to non-differentiable predictors using tools from [Lyu and Li, 2020] is interesting, and we will comment on this as a future direction. \n\n---\n\n**Q2**. By comparing the proofs of this work with that of https://arxiv.org/pdf/2402.15926 , it seems like similar proof technique can be utilized to incorporate the three assumptions to derive the same results involving the three constants.\n\n**A2**. We emphasize that our results are for nonlinear networks, but [Wu et al., 2024] only focused on linear predictors (or networks in the NTK regime). While our EoS analysis is motivated by [Wu et al., 2024], identifying the correct set of assumptions to enable the analysis for nonlinear networks requires nontrivial efforts. \n\nBesides, our stable phase analysis also proves margin improvement for non-homogenous predictors, which partially solves an open problem noted by [Ji and Telgarsky, 2020]. To prove this result, we use techniques significantly different from those of [Wu et al., 2024]. \n\n\n---\n\n**Q3**. Will the loss landscape with non-homogenous layers be the same as the linear networks case? For example, the authors show that in logistic regression, the landscape is a valley instead of a quadratic basin. The iterates converges quickly in max-margin subspace direction and the oscillations within the orthogonal direction becomes small progressively with time. Can you give an intuition of how the loss landscape may look for non-homogenous layers?\n\n**A3**. No, the loss landscape for non-homogeneous networks is significantly different from that of logistic regression. Note that the formal is non-convex while the latter is convex. It is difficult to develop geometric intuitions about the non-convex landscape of non-homogeneous networks. This difficulty also demonstrates the significance of our contributions to establishing margin improvement and EoS theory for non-homogeneous networks.  \n\n---\n\n**Q4**. Is it true that unllike MSE losses, where there is a chance of divergence for very high lrs, the logistic regression losses always converge irrespective of how large lr is used.\n\n**A4**. This is a good question and we do not know the answer. In practice, the loss can diverge when training deep networks with very large stepsizes, even under logistic loss. However, we are not sure if this is a numerical issue or if there is a provable upper bound for the convergent stepsize. \n\n---\n\n**Q5**. I missed the part in the paper where the authors discuss progressive sharpening. Usually for MSE losses, PS first takes place unless the sharpness hits 2/lr after which loss oscillates around walls of basin. Does logistic regression loss do not exhibit PS? Seems like EOS starts from the offset and then oscillations stop when iterates enter a valley of flatter oscillation directions.\n\n**A5**. Empirically, progressive sharpening also happens under logistic loss, as shown by [Cohen et al., 2020]. However, this is beyond the scope of this work, which focuses on the implicit bias and convergence of large stepsize GD. \n\n---", "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": "2Q4aonaT6B", "reviewer": "Reviewer_cshA", "summary": "This paper studies the effect of large learning rates for near-homogeneous logistic regressions. Especially, it extends previous results on linear predictors under large learning rates (an EoS phase and a stable phase for the convergence of GD; faster convergence) to nonlinear predictors with Lipschitzness and Lipschitz smoothness, for example, the two-layer neural networks with last layer fixed. The authors also prove the margin improvement for large learning rates and this near-homogeneous model. In the end, they demonstrate a faster convergence using large learning rates.", "questions": "1. What is the initial condition in Theorem 3.2?\n2. For the margin improvement part (Theorem 2.2), it would be clearer if the dependence on the learning rates in the modified margin function could be explained in the main body. Also, it would be good if there were direct comparisons between small and large learning rates results of margin improvement, e.g., different dependence on learning rates or other parameters. This may help better illustrate the advantages of using large learning rates.", "limitations": "Although the authors improve the theory from linear models to near-homogeneous models, it is still far from the ones used in practice. If possible, some discussions on the intuitions of more nontrivial nonlinear models may be helpful.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This work is the first to prove the margin improvement for large learning rates and this near-homogeneous model.\n2. The authors analyze a nonlinear (near-homogeneous) model, which is closer to practice than previous works.", "weaknesses": "The role and effect of large learning rates, especially compared to small learning rates, may need more illustration.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your positive feedback. We answer your questions as follows. \n\n---\n\n**Q1**. The role and effect of large learning rates, especially compared to small learning rates, may need more illustration.\n\n**A1**. Corollary 4.2 and Theorem 4.3 together show a separation between large and small stepsize, where a large stepsize enables $O(1/T^2)$ loss, but small stepsizes suffer from at least a $\\Omega(1/T)$ loss. This illustrates the effect of large stepsize for fast optimization. \n\n---\n\n**Q2**. What is the initial condition in Theorem 3.2?\n\n**A2**. We do not have explicit assumptions on the initial condition of $w_0$. For the bound in Theorem 3.2 to be non-vacuous, $|| w_0 ||$ needs to be small compared to $\\eta t$. \n\n---\n\n**Q3**.  For the margin improvement part (Theorem 2.2), it would be clearer if the dependence on the learning rates in the modified margin function could be explained in the main body. Also, it would be good if there were direct comparisons between small and large learning rates results of margin improvement, e.g., different dependence on learning rates or other parameters. This may help better illustrate the advantages of using large learning rates.\n\n**A3**. The modified margin is defined in Equation (10) in the appendix with explicit dependence on the stepsize. We will add this to the main paper. \n\nThe impact of the stepsize on the margin improvement is complex. Let us use $s$ to denote the starting time of the stable phase. If we assume $w_s$ is independent of the stepsize, then the loss and the parameter norm depend on the stepsize through the bounds $\\Theta(1/\\eta(t-s))$ and $\\Theta(\\log(\\eta(t-s)))$, respectively. However, $w_s$ depends on the stepsize through a complex function, making it difficult to determine the quantitative effect of stepsize on margin improvement. \n\n---\n\n**A4**. Although the authors improve the theory from linear models to near-homogeneous models, it is still far from the ones used in practice. If possible, some discussions on the intuitions of more nontrivial nonlinear models may be helpful.\n\n**Q4**. We believe our intuitions about the transition from an EoS phase to a stable phase extend to more general nonlinear models. Specifically, the initial EoS phase happens when GD oscillates within a sharp valley, and GD enters the stable phase when it navigates itself into a flat valley. Our theory of large stepsize is consistent with the celebrated flat minima intuition. We will add this discussion in the revision. \n\n---", "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:
Thank you for all your feedbacks. We add two figures in the pdf. 

- Figure 1 shows the test accuracy of two-layer networks for CIFAR-10 under the same setting of Figure 2(d)-(f). The
results support the intuition that large stepsizes lead to stronger implicit biases with “nicer“ features.
- Figure 2 shows training loss and margins of a two-layer network with leaky softplus activations
on a synthetic linear separable dataset. There are  five samples in the dataset, which are $$((0.05, 1,2),1), ((0.05, -2,1),1), ((-1,0,2),-1), ((0.05,-2,-2),1),((0.05, 1,-2),1).$$ The max-margin direction is $(1,0,0)$ with a normalized margin of $0.05$.  The network only has two neurons with fixed weights $1/2$ and $-1/2$. The leaky softplus activation is $\tilde \phi(x) = ( x +  \phi(x))/2$, where $\phi$ is the softplus activation. The stepsize is $3$.  We can observe that both neurons have negative normalized margins during the training, while the network's normalized margin increases and becomes positive.

Meta-review: The paper received positive reviews, with a consensus among reviewers that this is a solid extension on recent work on neural network GD dynamics. Consequently, I recommend acceptance of this paper. When preparing the camera-ready revision, please pay close attention to the reviewers’ feedback and include the additional explanations provided during the rebuttal period.
Decision: Accept (poster)