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 l bartlett
u c berkeley
lawrence berkeley national laboratory
google  deepmind
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 09:59:20

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.

Main paper: 
 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 a few exceptions will be discussed later. However, in practice, GD is often used with a 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, it is 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 is relatively well understood. For instance, classical optimization theory suggests that GD minimizes convex and L-smooth functions if the stepsize η is well below 2/L, with a convergence rate of O(1/(ηt)), where t is 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-margin 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. Figure 1 : The behavior of (GD) for optimizing a non-homogenous four-layer MLP with GELU activation function on a subset of CIFAR-10 dataset. We randomly sample 6, 000 data with labels "airplane" and "automobile" from CIFAR-10 dataset. The normalized margin is defined as min i∈[n] y i f (w t ; x i )/∥w t ∥ 4 , which is close to (3). The blue curves correspond to GD with a large stepsize η = 0.2, where the empirical risk oscillates in the first phase but decreases monotonically in the second phase. The orange curves correspond to GD with a small stepsize η = 0.005, where the empirical risk decreases monotonically. Furthermore, Figure 1 (b) suggests the normalized margins of both two curves increase and converge in the stable phases. Finally, Figure 1 (c) suggests that large stepsize achieves a better test accuracy, consistent with larger-scale learning experiment [Hoffer et al., 2017 , Goyal et al., 2017] . More details can be found in Section 5.More recently, large stepsize GD that induces oscillatory risk has been analyzed in simplified setups [see Ahn et al., 2023 , Zhu et al., 2022 , Kreisler et al., 2023 , Chen and Bruna, 2023 , Wang et al., 2022a , Wu et al., 2023 , 2024 , for an incomplete list of references]. In particular, in logistic regression with linearly separable data, Wu et 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] but also for an arbitrarily large stepsize. In the same problem, Wu et al. [2024] further showed that large stepsize GD that undergoes risk oscillation can achieve an Õ(1/t 2 ) empirical risk, whereas small stepsize GD that monotonically decreases the empirical risk must suffer from a Ω(1/t) empirical risk. Nonetheless, these theories of large stepsize GD are limited to 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 a binary classification dataset (x i , y i ) n i=1 , where x i ∈ R d is a feature vector and y i ∈ {±1} is a binary label. For simplicity, we assume ∥x i ∥ ≤ 1 for all i throughout the paper. For a predictor f , the empirical risk under logistic loss is defined asL(w) := 1 n n i=1ℓ(y i f (w; x i )), ℓ(t) := log(1 + e -t ).(1)Here, the predictor f (w; •) : R d → R is parameterized by trainable parameters w and is assumed to be continuously differentiable with respect to w. The predictor is initialized from w 0 and then trained by gradient descent (GD) with a constant stepsize η > 0, that is, If κ = 0, then Assumption 1C requires the predictor to be exactly 1-homogenous. Our Assumption 1C allows the predictor to have a bounded homogenous error. It is clear that linear predictors f (w; x) := w ⊤ x satisfy Assumption 1 with ρ = sup i ∥x i ∥ ≤ 1, β = 0, and κ = 0. Another notable example is two-layer networks given by f (w; x) := 1 m m j=1 a j ϕ(x ⊤ w (j) ), w (j) ∈ R d , j = 1, . . . , m,where we assume a j ∈ {±1} are fixed and w ∈ R md , the stack of (w (j) ) m j=1 , 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 1 covers two-layer networks with many commonly used activations ϕ(•). The proof is provided in Appendix F.1.Example 2.1 (Two-layer networks). Two-layer networks defined in (2) with the following activation functions satisfy Assumption 1 with the described constants:• GELU. ϕ(x) := x 2 erf 1 + (x/ √ 2) with κ = e -1/2 / √ 2π, β = 2/m, and ρ = ( √ 2π + e -1/2 )/ √ 2πm.• Softplus. ϕ(x) := log(1 + e x ) with κ = log 2, β = 1/m, and ρ = 1/ √ m.• SiLU. ϕ(x) := x/(1 + e -x ) with κ = 1, β = 4/m, and ρ = 2/ √ m.• Huberized ReLU [Chatterji et al., 2021] . For a fixed h > 0,ϕ(x) :=      0 x < 0, x 2 2h 0 ≤ x ≤ h, x -h 2 x > h,with κ = h/2, β = 1/(hm), and ρ = 1/ √ m.Margin for nearly homogenous predictors. For a nearly homogenous predictor f (w; •) (see Assumption 1C), we define its normalized margin (or margin for simplicity) asEQUATIONA 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-homogenous predictors that satisfy Assumption 1C. The next theorem gives sufficient conditions for large stepsize GD to 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 to Appendix A. Theorem 2.2 (Stable phase and margin improvement). Consider (GD) with stepsize η on a predictor f (w; x) that satisfies Assumptions 1A and 1B. If there exists r ≥ 0 such thatEQUATION)then GD is in the stable phase for t ≥ r, that is, (L(w t )) t≥r decreases monotonically. If additionally the predictor satisfies Assumption 1C and there exists s ≥ 0 such thatEQUATIONwe have the following for t ≥ s:• Risk convergence. L(w t ) = Θ(1/t).• Parameter increase. ∥w t+1 ∥ ≥ ∥w t ∥ and ∥w t ∥ = Θ(log(t)).• Margin improvement. There exists a modified margin function γ c (w) such that γ c (w t ) is increasing and bounded.γ c (w t ) is a multiplicative approximiator of γ(w t ), that is, there exists c > 0 such thatγ c (w t ) ≤ γ(w t ) ≤ 1 + c log(1/L(w t )) γ c (w t ), t ≥ s.As a direct consequence, lim t→∞ γ(w t ) = lim t→∞ γ c (w t ).Theorem 2.2 shows that for an arbitrarily large stepsize η, GD must enter the stable phase if the empirical risk falls below a threshold depending on η 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 a Θ(1/t) 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-homogenous.Limitations. The stable phase conditions in Theorem 2.2 require GD to enter a sublevel set of the empirical risk. However, such a sublevel set might be empty. For instance, let f (w; x) be a two-layer network (2) with sigmoid activation. Notice that the predictor is uniformly bounded, |f (w; x)| ≤ 1, so we haveL(w) = 1 n n i=1 log(1 + e -yif (w;xi) ) ≥ log(1 + e -1).On the other hand, we can also verify that Assumption 1C is satisfied by f (w; x) with κ = 1 but no smaller κ. Therefore (5) cannot be satisfied. In general, the sublevel set given by the right-hand side of Assumption 1C is non-empty if there exists a unit vector v such that mini y i f (λv; x i ) → ∞ as λ → ∞.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 can be modified by adding a 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 to 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 is 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-homogenous 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. 
edge of stability phase: 
 Our stable phase results in Theorem 2.2 require the risk to be below a certain threshold (see ( 4) and ( 5)). In 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 in 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 ϕ : R → R be continuously differentiable. Moreover, A. Derivative condition. Assume there exists 0< α < 1 such that α ≤ |ϕ ′ (z)| ≤ 1. B. Smoothness. Assume there exists β > 0 such that for all x, y ∈ R, |ϕ ′ (x) -ϕ ′ (y)| ≤ β|x -y|.C. Near homogeneity. Assume there exists κ > 0 such that for every z ∈ R, |ϕ(z) -ϕ ′ (z)z| ≤ κ.Recall that sup i ∥x i ∥ ≤ 1. One can then check by direct computation that, under Assumption 2, two-layer networks (2) satisfy Assumption 1 with ρ = 1/ √ m, β = β/m, and κ = κ.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 generalized to allow |ϕ ′ (z)| ≤ C for a constant C > 0. The other condition in Assumption 2A, |ϕ ′ (z)| ≥ α, however, is non-trivial. This condition is widely used in 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 ≤ c < 1.• Let ϕ be GELU, Softplus, or SilU in Example 2.1, then its modification φ(x) := cx+(1-c)/4•ϕ(x) satisfies Assumption 2 with κ = 1, α = 0.25, and β = 1. In particular, the modification of softplus can be viewed as a smoothed leaky ReLU.• Let ϕ be the Huberized ReLU in Example 2.1, then its modification φ(x) := cx + (1 -c)/4 • ϕ(x)satisfies Assumption 2 with κ = h/2, α = 0.5, and β = 1/4h.• The "leaky" tanh, φ(x) := cx + (1 -c) tanh(x), and the "leaky" sigmoid, φ(x) := cx + c/(1 + e -x ), both satisfy Assumption 2 with κ = 1, α = 0.5 and β = 1.For the technical difficulty of non-convex optimization, we also need to assume a linearly separable dataset to conduct our EoS phase analysis. Assumption 3 (Linear separability). Assume there is a margin γ > 0 and a unit vector w * such that y i x ⊤ i w * ≥ γ for every i = 1, . . . , n.Assumption 3 serves as a sufficient condition for two-layer neural networks, regardless of width, to reach the initial bound of 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 η := mη, where m is the network width and η can be arbitrarily large. Then for every t > 0, we have1 t t-1 k=0 L(w k ) ≤ 1 + 8 log 2 (γ 2 ηt)/α 2 + 8κ 2 /α 2 + η 2 γ 2 ηt + ∥w 0 ∥ 2 mηt = O log 2 (ηt) + η 2 ηt .Theorem 3.2 suggests that the average risk of training two-layer networks decreases even when GD is used with large stepsize. Consequently, the risk thresholds (4) and ( 5) for GD to 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 1 in [Wu et al., 2024] , Theorem 3.2 extends their EoS phase bound from linear predictors to two-layer networks. 
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, a direct application of Theorem 3.2 only leads to a 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 1 with ρ = 1/ √ m, β = β/m, and κ = κ. Denote the stepsize by η := mη, where m is the network width and η > 0 can be arbitrarily large.• Phase transition time. There exists s ≤ τ such that (5) in Theorem 2.2 holds, whereτ := 128(1 + 4κ) α 2 max c 1 η, c 2 n, e, c 2 η + c 1 n η log c 2 η + c 1 n η , (c 2 η + c 1 n) η • ∥w 0 ∥ √ m ,where c 1 := 4e κ+2 and c 2 := (8 + 4 β). Therefore (GD) is in the stable phase from s onwards.• Explicit risk bound in the stable phase. We have (L(w t )) t≥s monotonically decreases andL(w t ) ≤ 2 α 2 γ 2 η(t -s) , t ≥ s.Theorems 2.2, 3.2 and 4.1 together characterize the behaviors of large stepsize GD in training twolayer 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 a certain stepsize-dependent 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 a flatter valley. We believe this insight generalizes to 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 a faster optimization in the stable phase but stays longer in the EoS phase. Choosing a suitably large stepsize that balances the steps in 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 thatT ≥ 256(1 + 4κ) α 2 γ 2 max c 1 n, 4c 2 2 , 2c 2 ∥w 0 ∥ √ m ,where c 1 := 4e κ+2 and c 2 := (8 + 4 β) are as in Theorem 4.1. Then for stepsize η := ηm, whereη := α 2 γ 2 256(1 + 4κ)c 2 T,we have τ ≤ T /2 andL(w T ) ≤ 2048(1 + 4κ)c 2 α 4 γ 4 • 1 T 2 = O 1 T 2 .Theorem 1 = (γ, 1 -γ 2 ), x 2 = (γ, -1 -γ 2 /2), y 1 = y 2 = 1, 0 < γ < 0.1.It is clear that (x i , y i ) i=1,2 satisfy Assumption 3. If (L (w t )) t≥0 is non-increasing, thenL (w t ) ≥ c 0 /t, t ≥ 1where c 0 > 0 is a function of (α, ϕ, x 1 , x 2 , γ, κ, β) but is independent of t and η.Effect of model rescaling. We conclude this section by discussing the impact of rescaling the model. Specifically, we replace the two-layer network in the mean-field scaling (2) by the followingf (w; x) := b • 1 m m j=1 a j ϕ(x ⊤ w (j) ),and evaluate the impact of the scaling factor b on our 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 a fixed budget of T steps and get the following bound. Detailed derivations are deferred to Appendix D.L(w T ) = O(1/T 2 ) if b ≥ 1, O(b -3 /T 2 ) if b < 1.This suggests that as long as b ≥ 1, we get the same acceleration effect. In particular, the mean-field scaling b = 1 [Song et al., 2018, Chizat and Bach, 2020] and the neural tangent kernel et al., 2018 et al., , 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.(NTK) scaling b = √ m [Du 
experiments: 
 We conduct three sets of experiments to validate our theoretical insights. In the first set, we use a 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 a multilayer perceptron (MLP) with four trainable layers and GELU activation functions, with a hidden dimension of 200 for each hidden layer. The MLP is trained using gradient descent with random initialization, as described in (GD). The results are shown in Figures 1(a ) to 1(c).In the second set of experiments, we consider an XOR dataset consisting of four samples:x 1 = (-1, -1), y 1 = 1; x 2 = (1, 1), y 2 = 1; x 3 = (1, -1), y 3 = -1; x 4 = (-1, 1), y 4 = -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 c = 0.5). The network width is m = 20. The initialization is 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 on a 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(f).Margin improvement. Figures 1(b), 2(c) and 2(f) show that the normalized margin nearly monotonically 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 1(a) , 2(a) and 2(d), we observe that after GD enters the stable phase, a larger stepsize consistently leads to a smaller empirical risk compared to the smaller stepsizes, which is consistent with our Theorem 4.1 and Corollary 4.2. Besides, Figures 2(b ) and 2(e) suggest that, asymptotically, GD converges at a rate of O(1/(ηt)) = O(1/(ηt)) (The width of networks is fixed), which verifies the sharpness of our stable phase bound in Theorem 4.1.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. A detailed example illustrating this is provided in Appendix E. 
related works: 
 In this section, we discuss related papers. We consider an XOR dataset and a subset of CIFAR-10 dataset. In both cases, we observe that (1) GD with a large stepsize achieves a faster optimization compared to GD with a small stepsize, (2) the asymptotic convergence rate of the empirical risk is O(1/(ηt)), 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 Gunasekar et al. [ , 2018]] [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 a large stepsize that induces non-monotonic risk.Two papers [Nacson et al., 2019a , Kunin et al., 2022] studied margin maximization theory for a special form of non-homogenous models. Specifically, when viewed in terms of different subsets of the trainable parameters, the model is homogeneous, although the order of homogeneity may vary. Compared to their setting, our non-homogenous models only require a bounded homogenous error (see Assumption 1C). Therefore, our theory can 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 to 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 oneor two-dimensional functions [Zhu et al., 2022 , Chen and Bruna, 2023 , Ahn et al., 2022 , Kreisler et al., 2023 , Wang et al., 2023] , linear model [Wu et al., 2023 [Wu et al., , 2024]] , matrix factorization [Wang et al., 2022a, Chen and Bruna, 2023] , scale-invariant networks [Lyu et al., 2022] , linear networks under MSE loss [Ren et al., 2024 , Andriushchenko et al., 2023] , for an incomplete list of references.Compared to them, we focus on a 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 to ours.In what follows, we make a detailed discussion about papers that directly motivate our work [Lyu and Li, 2020 , Ji and Telgarsky, 2020 , Chatterji et al., 2021 , Wu et 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 {w t /∥w t ∥} ∞ t=1 converging to KKT points of a margin-maximization 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] , Ji 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 1-homogenity condition to a "near L-homogeneity" condition for a general L. Chatterji et al. [2021] . The work by Chatterji et al. [2021] studies the convergence of GD in training deep networks under logistic loss. Their results are related to ours as we both consider networks with nearly homogeneous activations and we both have a stable phase analysis (although this is not explicitly mentioned in their paper). However, our results are significantly different from theirs. Specifically, in our notation, they require the homogenous error κ (see Assumption 2C) to be smaller than O(log(1/L(w s ))/∥w s ∥) ≈ O(γ(w s )), where s is the time for GD to enter the stable phase. Note that the margin when GD enters the stable phase could be arbitrarily small. In comparison, we only require the homogenous error to 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 a small h in Example 2.1. Moreover, they require the stepsize η to be smaller than O(κ/∥w s ∥ 8 ) , thus they only allow very small stepsize. In contrast, we allow η to be arbitrarily large.Comparison with Wu et al. [2023 Wu et al. [ , 2024]] . The works by Wu et al. [2023 Wu et al. [ , 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 [Wu et al., , 2024] ] from linear predictors to non-linear predictors such as 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. 
conclusion: 
 We provide a 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-homogenous predictors. Finally, we show that large stepsize GD, by entering the initial oscillatory phase, achieves acceleration when minimizing the empirical risk.Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMS-2031883 and #814639, and of the ONR through MURI award N000142112431. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: 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 a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We 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 or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • 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.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing 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 as 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
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 an equivalent approval/review based on the requirements of your country or institution) were obtained?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. • Depending on 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.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In this section, we will prove results for a general smooth predictor f (w; 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.• q i (t) := y i f (w t ; x i ), q min (t) := min i∈[n] q i (t).• L t := L(w t ), ρ t := ∥w t ∥ 2 .Then, we have the following expression:L(w t ) = 1 n n i=1 ℓ(q i (t)).Here, we give a summary of this section. The proofs are organized into 5 parts.• In Appendix A.1, we characterize the decrease of loss L t .• In Appendix A.2, we characterize the change of the parameter norm ρ t .• In Appendix A.3, we show the convergence of the normalized margin γ(w t ).• In Appendix A.4, we characterize the sharp rates of loss L t and parameter norm ρ t .• In Appendix A.5, we give the proof of Theorem 2.2.In this section, we will show that the loss L t decreases monotonically in the stable phase. To begin with, we introduce the following definition which is another characterization of β-smoothness. Definition 1 (Linearization error). Given a continuously differentiable function f : R d → R and two points w, v ∈ R d , the linearization error of f (v) with respect to w is:ξ[f ](w, v) := f (v) -f (w) -∇f (w) ⊤ (v -w).For a β-smooth function, standard convex optimization theory gives the following linearization error bound.Fact A.1 (Linearization error of β-smooth function). For a β-smooth function f : R d → R, we haveξ[f ](w, v) := f (v) -f (w) -∇f (w) ⊤ (v -w) ≤ β 2 ∥v -w∥ 2 2, for every w and v.We first show a 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 f to be twice differentiable, extra efforts are needed. Lemma A.2 (Self-boundedness of logistic loss). For the logistic loss ℓ(z):= log(1 + exp(-z)), we have 0 ≤ ℓ(z) -ℓ(x) -ℓ ′ (x)(z -x) ≤ 2ℓ(x)(z -x) 2 for |z -x| < 1.Proof of Lemma A.2. See the proof of Proposition 5 in [Wu et al., 2024] . The lower bound is by the convexity of ℓ(•).The next lemma controls the decrease of the risk L t . Lemma A.3 (Decrease of L t ). Suppose Assumptions 1A and 1B hold. If L(w t ) ≤ 1 ηρ 2 , then we have-η(1 + β ηL(w t ))∥∇L(w t )∥ 2 ≤ L(w t+1 ) -L(w t ) ≤ -η(1 -(2ρ 2 + β)ηL(w t ))∥∇L(w t )∥ 2 .Particularly, this indicates that ifL(w t ) ≤ 1 η(2ρ 2 +β) , then L(w t+1 ) ≤ L(w t ).Proof of Lemma A.3. By Assumptions 1A and 1B, we have ∥∇f ∥ 2 ≤ ρ and f (w; x) is β-smooth as a function of w. Therefore, for every i ∈ [n] we have|q i (t + 1) -q i (t)| = |y i (f (w t+1 ; x i ) -f (w t ; x i ))| = ∇f (w t + θ (w t+1 -w t ) ; x i ) ⊤ (w t+1 -w t ) by intermediate value theorem ≤ ρ ∥w t+1 -w t ∥ ≤ ρη∥∇L t ∥ since w t+1 = w t -η∇L t ≤ ρ 2 ηL t ≤ 1. since ∥∇L t ∥ ≤ L t ρ Then by Lemma A.2, we have ℓ(q i (t + 1)) ≤ ℓ(q i (t)) + ℓ ′ (q i (t))(q i (t + 1) -q i (t)) + 2ℓ(q i (t))(q i (t + 1) -q i (t)) 2 ≤ ℓ(q i (t)) + ℓ ′ (q i (t))⟨y i ∇f (w t ; x i ), w t+1 -w t ⟩ + |ℓ ′ (q i (t))| • |ξ[f ](w t , w t+1 )| + 2ℓ(q i (t))(q i (t + 1) -q i (t)) 2 since q i (t + 1) -q i (t) = ⟨y i ∇f (w t ; x i ), w t+1 -w t ⟩ + y i ξ[f ](w t , w t+1 ) ≤ ℓ(q i (t)) + ℓ ′ (q i (t))⟨y i ∇f (w t ; x i ), w t+1 -w t ⟩ + ℓ(q i (t))(β + 2ρ 2 )∥w t+1 -w t ∥ 2 .by Fact A.1 and the previous inequality Taking an average over all data points, we haveL t+1 ≤ L t -η∥∇L t ∥ 2 + (2ρ 2 + β)η 2 L t ∥∇L t ∥ 2 , which is equivalent to L t+1 -L t ≤ -η(1 -(2ρ 2 + β)ηL t )∥∇L t ∥ 2 .We complete the proof of the right hand side inequality. The left hand side inequality can be proved similarly. In detail, we can show that:ℓ(q i (t + 1)) ≥ ℓ(q i (t)) + ℓ ′ (q i (t))(q i (t + 1) -q i (t) ≥ ℓ(q i (t)) + ℓ ′ (q i (t))⟨y i ∇f (w t ; x i ), w t+1 -w t ⟩ -|ℓ ′ (q i (t))| • |ξ[f ](w t , w t+1 )|.Taking the average over all data points, we have L t+1 ≥ L t -η(1 + β ηL t )∥∇L t ∥ 2 . Now we have completed the proof.In this section, we demonstrate that the parameter norm, ρ t , increases monotonically during the stable phase. We introduce a crucial quantity, v t , defined as the inner product of the gradient and the negative weight vector: v t := ⟨∇L(w t ), -w t ⟩. This quantity, v t , plays a key role in controlling the increase of the parameter norm. Notably, v t appears as the cross term in the expression ∥w t+1 ∥ 2 = ∥w t -η∇L(w t )∥ 2 . By managing v t , we can effectively characterize the increase in the parameter norm.Recall that our loss function is ℓ(x) := log(1 + e -x ). Inspired by Lyu and Li [2020] , we define the following two auxiliary functions for the logistic loss:EQUATION)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 ψ(x) and ι(x) are close to the identity function, i.e., ψ(x) ≈ ι(x) ≈ x, for x large enough. Then, we have an exponential-loss-like decomposition of L t :EQUATIONThese two functions ψ, ι will help us to analyze the lower bound of v t . First, we list some properties of ψ and ι here.Lemma A.4 (Auxiliary functions of ℓ). The following claims hold for ℓ, ψ, and ι.• ℓ(x) = e -ψ(x) .• ℓ is monotonically decreasing, while ψ and ι are monotonically increasing.• ψ ′ (ι(x)) = 1 ι ′ (x) ; • ψ ′ (x)x is increasing for x ∈ (0, +∞).Proof of Lemma A.4. The first two properties are straightforward. For the third property, we apply chain rule on ψ(ι(x)) = x to getψ ′ (ι(x))ι ′ (x) = 1.For the fourth property, notice thatψ ′ (x)x = x (1 + e x ) log(1 + e -x ).The denominator is positive and decreasing sinced dx (1 + e x ) log(1 + e -x ) = e x log(1 + e -x ) -1 ≤ e x e -x -1 = 0.Combining this with the fact that x is positive and increasing, we have the desired result.Besides, we have the following property of ι. This is the key lemma to handle the homogeneous error.Actually, this lemma is another way to show ι(x) is close to the identity function.Lemma A.5 (Property of ι). For every x ∈ R, we haveι(x) ι ′ (x) ≥ x + log log 2.Proof of Lemma A.5. Recall that ι(x) = -log(e e -x -1), ι ′ (x) = e e -x e -x e e -x -1 .Let y = e -x . We haveι(x) ι ′ (x) =-log(e e -x -1)(e e -x -1) e e -x e -x = -log(e y -1)(e y -1) e y y .Define s(y) := ι(x) ι ′ (x) -x -log log 2. Then, s(y) = -log(e y -1)(e y -1) e y y + log(y) -log log 2,s ′ (y) = log(e y -1) • e y -y -1 e y y 2 >0 .Note that the sign of s ′ is determined by log(e y -1). For 0 < e y ≤ 2, s ′ (y) ≤ 0 and s(y) is decreasing; for e y ≥ 2, s(y) is increasing. Therefore,min y∈(0,∞) s(y) = s(log 2) = 0.Since x = -log y, we have the desired result.Another important property of ι is that it can provide a lower bound for q min (t). Lemma A.6 (ι bound q min ). For every t ≥ 0, we haveq min (t) ≥ ι -log(L t ) -log n .Proof of Lemma A.6. We use (7) to get1 n ℓ(q min (t)) ≤ L t ⇒ 1 n e -ψ(qmin(t)) ≤ L t ⇒ ψ(q min (t)) ≥ -log n -log L t ⇒ q min (t) ≥ ι -log(L t ) -log n . by Lemma A.4Then, we complete the proof. Now, we are ready to give a lower bound of v t . The following lemma is an extension of Corollary E.6 in Lyu and Li [2020] , where they dealt with a homogeneous model and the exponential loss; we extend this to a non-homogeneous model. The key ingredient is Lemma A.5. Lemma A.7 (A lower bound of v t ). Suppose Assumption 1C holds. Consider v t := ⟨∇L(w t ), -w t ⟩.If L t ≤ 1 2ne κ , then v t ≥ -L t log(2ne κ L t ) ≥ 0.Proof of Lemma A.7. By definition, we havev t := ⟨∇L(w t ), -w t ⟩ = - 1 n n i=1 ℓ ′ (y i f (w t ; x i ))y i ⟨∇f (w t ; x i ), w t ⟩ = - 1 n n i=1 ℓ ′ (y i f (w t ; x i ))y i f (w t ; x i ) - 1 n n i=1 ℓ ′ (y i f (w t ; x i )) y i ⟨∇f (w t ; x i ), w t ⟩ -y i f (w t ; x i ) ≥ - 1 n n i=1 ℓ ′ (y i f (w t ; x i ))y i f (w t ; x i ) -κL t since |ℓ ′ (x)| ≤ ℓ(x) and |⟨∇f (w t ; x i ), w t ⟩ -f (w t ; x i )| ≤ κ by Assumption 1C = 1 n n i=1e -ψ(qi(t)) ψ ′ (q i (t))q i (t) -κL t . since ℓ(•) = exp(-ψ(•))Applying Lemma A.6 and Lemma A.4, we have q i (t) ≥ q min (t) ≥ ι -log(nL t ) := -log(e nLt -1) ≥ -log(e 1 2 -1) ≥ 0.Then we can apply Lemma A.4 to getψ ′ (q i (t))q i (t) ≥ ψ ′ ι -log(nL t ) ι -log(nL t ) = ι -log(nL t ) ι ′ -log(nL t ) .Invoking Lemma A.5, we haveι -log(nL t ) ι ′ -log(nL t ) ≥ -log(nL t ) + log log 2 ≥ -log(nL t ) + log log e 1 2 = -log(2nL t ).Putting the above two inequalities together, we haveψ ′ (q i (t))q i (t) ≥ -log(2nL t ), for every i = 1, . . . , n.Plugging this back to the bound of v t , we getv t ≥ - 1 n n i=1 e -ψ(qi(t)) log(2nL t ) -κL t = -L t log(2nL t ) -κL t = -L t log(2ne κ L t ) ≥ 0.This completes the proof.Right now, we get a lower bound for v t , which is the cross term in the expression of ∥w t+1 ∥ 2 . The next lemma controls the increase of the parameter norm ρ t using v t and L t .Lemma A.8 (The increase of ρ t ). Suppose Assumptions 1A and 1C hold.If L t ≤ min 1 2ne κ , 1 η(4ρ 2 +2β) , then 0 ≤ 2ηv t ≤ ρ 2 t+1 -ρ 2 t ≤ 2ηv t • 1 - 1 2 log(2ne κ L t ).Proof of Lemma A.8. By definition, we haveρ 2 t+1 -ρ 2 t = 2η⟨∇L t , -w t ⟩ + η2 ∥∇L t ∥ 2 = 2ηv t + η2 ∥∇L t ∥ 2 ≥ 2ηv t ≥ 0,where the last inequality is by Lemma A.7. Besides,ρ 2 t+1 -ρ 2 t = 2ηv t 1 + η∥∇L t ∥ 2 2v t ≤ 2ηv t 1 + ηL 2 t ρ 2 2v tby ℓ ′ ≤ ℓ, Assumption 1A, and Lemma A.7≤ 2ηv t 1 + L t 2v t by L t ≤ 1 η(4ρ 2 + 2β) ≤ 2ηv t 1 - 1 2 log(2ne κ L t ). by Lemma A.7This completes the proof.In this section, we show that the normalized margin of a general predictor converges in the stable phase. Recall that we define the (normalized) margin asγ(w) := min i∈[n] y i f (w; x i ) ∥w∥ 2 .However, this normalized margin is not a smooth function of w. Instead, we consider a smoothed margin γ a as an easy-to-analyze approximator of the normalized margin [Lyu and Li, 2020 ]EQUATION)We see that γ a is a good approximator of γ. We can then use γ a to analyze the convergence of the normalized margin since they share the same limit (if it exists). While γ a is relatively easy to analyze for gradient flow [Lyu and Li, 2020] , analyzing that for GD with a large (but fixed) stepsize is 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 asEQUATIONand a modified margin asEQUATION)These two margins provide a second-order correction when viewing large stepsize GD as a first-order approximation of gradient flow. In the following discussion, we will show that γ(w) ≈ γ a (w) ≈ γ b (w) ≈ γ c (w). At last, we will use the convergence of γ c (w t ) to prove γ(w t ) converges.The following lemma shows that γ(w) ≈ γ a (w). Lemma A.9 (Smoothed margin). For the smooth margin γ a (w t ) defined in (8) and the normalized margin γ(w t ) defined in (3), we have• When L t ≤ 1 2n , we have q min (t) ≤ -log L t ≤ log(2n) + q min (t), and γ(w t ) ≤ γ a (w t ) ≤ γ(w t ) + log(2n) ρ t .• Assume Assumption 1C holds. If L t → 0, then |γ a (w t ) -γ(w t )| → 0.Proof of Lemma A.9. To prove the first claim, notice thatL t ≤ 1 2n =⇒ ℓ(q min (t)) = log(1 + exp(-q min (t))) ≤ nL t ≤ 1 2 .Therefore we have e -qmin(t) ≤ e t) ) ≤ e -qmin(t) .1 2 -1 ≤ 1. Using x 2 ≤ log(1 + x) ≤ x for 0 ≤ x ≤ 1, we get 1 2 e -qmin(t) ≤ ℓ(q min (t)) = log(1 + e -qmin(Then we can bound L t by1 2n e -qmin(t) ≤ 1 n ℓ(q min (t)) ≤ L t ≤ ℓ(q min (t)) ≤ e -qmin(t) ,which is equivalent to q min (t) ≤ -log L t ≤ log(2n) + q min (t).Dividing both sides by ρ t proves the second claim:γ(w t ) := q min (t) ρ t ≤ γ a (w t ) := -log L t ρ t ≤ γ(w t ) + log(2n) ρ t = log(2n) + q min (t) ρ t .For the last claim, we only need to show that ρ t → ∞. This is because if L t → 0, we have for anyi ∈ [n], y i f (w t ; x i ) → ∞. Using y i f (w t ; x i ) ≤ C r,κ ∥w t ∥ + C r from Lemma G.2, we have ρ t = ∥w t ∥ 2 → ∞.The following lemma shows that γ c (w) ≈ γ(w). Lemma A.10 (Modified and auxiliary margins). Suppose that Assumption 1 holds. For the modified margin γ c (w t ) defined in (10) and the auxiliary margin γ b (w t ) defined in (9), we have• If L t ≤ 1 2ne κ+2 , there exists a constant c such that γ c (w t ) ≤ γ b (w t ) ≤ γ(w t ) ≤ 1 + c log(1/L(w t )) γ c (w t ).Proof of Lemma A.10. To prove the first two inequalities, notice thate Φ(Lt) = -log(2ne κ L t ) • exp 1 + (4ρ 2 + 2β)η log(2ne κ L t ) using (10) ≤ -log(2ne κ L t ). since L t ≤ 1 2ne κ , exp 1 + (4ρ 2 + 2β)η log(2ne κ L t ) ≤ 1, and log(2ne κ L t ) > 0 ≤ -log(L t ) -log(2n) ≤ q min (t).By argument 1 in Lemma A.9Using the above, ( 8) to (10), we haveγ c (w t ) := e Φ(L(wt)) ∥w t ∥ ≤ -log(2ne κ L t ) ∥w t ∥ =: γ b (w t ) ≤ q min (t) ∥w t ∥ =: γ(w t ).Then, we will prove the remaining inequality. First, we haveγ(w t ) γ c (w t ) = γ(w t ) γ b (w t ) • γ b (w t ) γ c (w t ) = q min (t) -log(2ne κ L t ) • exp 1 + (4ρ 2 + 2β)η -log(2ne κ L t ) By the definitions of γ, γ b , γ c ≤ -log(L t ) -log(2ne κ L t ) • exp 1 + (4ρ 2 + 2β)η -log(2ne κ L t )Since q min (t) ≤ log(-L t ) by Lemma A.9= 1 + log(2ne κ ) -log(2ne κ L t ) • exp 1 + (4ρ 2 + 2β)η -log(2ne κ L t ) .To simplify the notation, we let c 1 := 1 + (4ρ2 + 2β)η and c 2 = log(2ne κ ).Since L t ≤ 1 2ne κ+2 ⇒ -log(2ne κ L t ) ≥ 2 > 1, we have 1 + (4ρ 2 + 2β)η -log(2ne κ L t ) = c 1 -log(2ne κ L t ) ≤ c 1 .Besides, given x < c, we have e x ≤ 1 + e c x. Therefore,exp 1 + (4ρ 2 + 2β)η -log(2ne κ L t ) = exp c 1 -log(2ne κ L t ) ≤ 1 + c 1 exp(c 1 ) -log(2ne κ L t ).Plugging this into the bound for γ(w t )/γ c (w t ), we getγ(w t ) γ c (w t ) = 1 + c 2 -log(2ne κ L t ) • exp c 1 -log(2ne κ L t ) ≤ 1 + c 2 -log(2ne κ L t ) • 1 + exp(c 1 )c 1 -log(2ne κ L t ) ≤ 1 + c 2 + exp(c 1 )c 1 + c 2 c 1 exp(c 1 ) -log(2ne κ L t ) Since -log(2ne κ L t ) ≥ 1 = 1 + c 2 + exp(c 1 )c 1 + c 2 c 1 exp(c 1 ) -log L t -c 2 . Note that -log L t -c 2 ≥ 2 > 1. Because x x-c2 is decreasing when x ≥ c 2 + 1, we have -log L t -log L t -c 2 ≤ c 2 + 1 =⇒ 1 -log L t -c 2 ≤ c 2 + 1 -log L t .Plug this inequality into the previous bound for γ(w t )/γ c (w t ) and we get γ(w t )γ c (w t ) ≤ 1 + (c 2 + exp(c 1 )c 1 + c 2 c 1 exp(c 1 ))(c 2 + 1) -log L t .Let c := (c 2 + exp(c 1 )c 1 + c 2 c 1 exp(c 1 ))(c 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 γ c (w t ) in the gradient descent dynamics. Specifically, we are going to use the property that Φ(x) -Φ(y) ≥ Φ ′ (y)(x -y), for all x, y. Lemma A.11 (Convexity of Φ). The function Φ(x) defined in (10) is convex for 0 < x < 1 2ne 2+κ .Proof of Lemma A.11. Check thatΦ ′ (x) = 1 - 1 log(2ne κ x) (1 + (4ρ 2 + 2β)η) x log(2ne κ x) ,and thatΦ ′′ (x) = (1 + (4ρ 2 + 2β)η) • (2 + log(2ne κ x)) -log 2 (2ne κ x) -log(2ne κ x)x 2 log 3 (2ne κ x) .Note that when x ≤ 1 2ne 2+κ , we have log(2ne κ x) ≤ -2, which implies Before we dive into the proof of the monotonic increasing γ c (w t ), we show that γ c is bounded first. The convergence of γ c is a direct consequence of the monotonic increasing and the boundedness of γ c .Lemma A.12 (An upper bound on γ c , γ b , γ a and γ c ).When L t ≤ 1 2ne κ , 1 η(4ρ 2 +2β) for t ≥ s, there exists B 0 such that γ c (w t ) ≤ γ b (w t ) ≤ γ a (w t ) ≤ γ(w t ) + log 2n ρ s ≤ B 0 .Proof. Apply lemma A.8, we have ∥w t ∥ ≥ ρ t ≥ ρ s . Then we can apply lemma G.2 and there exists a constant C ρs,κ such that for all i,|y i f (w t ; x i )| ≤ C ρs,κ ∥w t ∥.Hence,γ(w t ) = arg min i∈[n] y i f (w t ; x i ) ∥w t ∥ ≤ C ρs,κ .Besides, by Lemma A.9, we haveγ a (w t ) ≤ γ(w t ) + log 2n ρ t ≤ C ρs,κ + log 2n ρ s .By Lemma A.10, we haveγ c (w t ) ≤ γ b (w t ) ≤ γ a (w t ) ≤ C ρs,κ + log 2n ρ s .Let B 0 = C ρs,κ + log 2n ρs . 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θ t := w t ∥w t ∥ , ν t := θ t θ ⊤ t (-∇L t ), µ t := I -θ t θ ⊤ t (-∇L t ).Therefore, we have∥∇L t ∥ 2 = ∥ν t ∥ 2 + ∥µ t ∥ 2 .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 ν t is the component in the direction of w t , and the second component µ t is the component orthogonal to w t . We will show that the modified margin γ c (w t ) is monotonically increasing. And the increase of γ c (w t ) is lower bounded by a term that depends on ∥µ t ∥ 2 .Lemma A.13 (Modified margin is monotonically increasing). Suppose Assumption 1 holds. If there exists s such thatL s ≤ min 1 e κ+2 2n , 1 η(4ρ 2 + 2β) ,then for t ≥ s, we have• L t+1 ≤ L t . • v t ≥ -L t log(2ne κ L t ) ≥ 0. • ρ 2 t+1 -ρ 2 t ≥ 2ηv t . • log γ c (w t+1 ) -log γ c (w t ) ≥ ρ 2 t v 2 t ∥µ t ∥ 2 log ρt+1 ρt .As a consequence, γ c (w t ) admits a finite limit.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 haveL t+1 -L t := L(w t+1 ) -L(w t ) ≤ -η(1 -(2ρ 2 + β)ηL t )∥∇L t ∥ 2 .Multiplying both sides by 2v t 1-1 2 log(2ne κ L t ) 1-(2ρ 2 +β)ηLt > 0, we get 1 - 1 2 log(2ne κ Lt) 1 -(2ρ 2 + β)ηL t 2v t (L t+1 -L t ) ≤ -2ηv t 1 - 1 2 log(2ne κ L t ) ∥∇L t ∥ 2 .From Lemma A.8 we have0 ≤ ρ 2 t+1 -ρ 2 t ≤ 2ηv t 1 - 1 2 log(2ne κ L t ).Using the above we getEQUATION)Recall that ∥∇L t ∥ 2 = ∥ν t ∥ 2 + ∥µ t ∥ 2 . For ν t , we have∥ν t ∥ = 1 ρ t ⟨w t , -∇L t ⟩ = v t ρ t .Then we can decompose ∥∇L t ∥ 2 asEQUATION)Plugging this into (11) and dividing both two sides by 2v 2 t , we have 1 -1 2 log(2ne κ Lt) (1 -(2ρ 2 + β)ηL t )v t (L t+1 -L t ) ≤ - 1 ρ 2 t (ρ 2 t+1 -ρ 2 t ) 1 2 + ρ 2 t 2v 2 t ∥µ t ∥ 2 .By Lemma A.7, we have v t ≥ -L t log(2ne κ L t ). DefineΨ(x) := - 1 - 1 2 log(2ne κ x) (1 -(2ρ 2 + β)ηx)x log(2ne κ x) .Then we haveEQUATION)We are going to show that Ψ(x) ≤ -Φ ′ (x). Note that when 0 < x ≤ min 1 e κ+2 2n , 1 η(4ρ 2 +2β) , we have log(2ne κ x) < 0 and 1 -(2ρ 2 + β)ηx) ≥ 1 2 > 0. Therefore, we haveEQUATION)To get an upper bound of Ψ(x), we just need an upper bound of J. Let a := -1 2 log(2ne κ x) > 0 and b := (2ρ 2 + β)ηx ∈ (0, 1/2]. Then we invoke Lemma G.4 to getJ := 1 + a 1 -b ≤ 1 + 2a + 2b = 1 - 1 log(2ne κ x) + (4ρ 2 + 2β)ηx.Recall that x ≤ 1 e κ+2 2n ≤ 1 2ne κ and 2ne κ ≥ 1. Then we apply Lemma G.3 to get x ≤ -1 log(2ne κ x) . Plugging this into the bound of J, we getEQUATIONPlugging ( 15) into ( 14), we haveΨ(x) = J • -1 x log(2ne κ x) ≤ - 1 - 1 log(2ne κ x) (1 + (4ρ 2 + 2β)η) x log(2ne κ x) = -Φ ′ (x),which verifies that Ψ(x) ≤ -Φ ′ (x). By this and ( 13), we haveΦ ′ (L t )(L t+1 -L t ) + φ ′ (ρ 2 t )(ρ 2 t+1 -ρ 2 t ) 1 2 + ρ 2 t 2v 2 t ∥µ t ∥ 2 ≥ 0,where φ(x) = -log x = log(1/x). Recall that for 0 < x ≤ 1 2ne κ+2 , Φ(x) is convex by Lemma A.11. By convexity of φ and Φ, we haveΦ(L t+1 ) -Φ(L t ) + log 1 ρ 2 t+1 -log 1 ρ 2 t 1 2 + ρ 2 t 2v 2 t ∥µ t ∥ 2 ≥ 0.By the definition of γ c in ( 10), this can be rewritten aslog γ c (w t+1 ) -log γ c (w t ) = (Φ(L t+1 ) -Φ(L t )) + log 1 ρ t+1 -log 1 ρ t ≥ -log 1 ρ 2 t+1 -log 1 ρ 2 t ρ 2 t 2v 2 t ∥µ t ∥ 2 = ρ 2 t v 2 t ∥µ t ∥ 2 log ρ t+1 ρ t ≥ 0,where the last inequality is because of Lemma A.8. We have shown that γ c (w t ) is monotonically increasing. By Lemma A.12, γ c (w t ) is bounded. Therefore γ c (w t ) admits a finite limit. This completes the proof.Right now, we have already proved that γ c (w t ) is monotonically increasing and bounded, which indicates γ c (w t ) converges. However, if we want to show that γ(w t ) converges, we still need to verify that L t → 0, which is the crucial condition for γ c (w t ), γ b (w t ), γ a (w t ), and γ(w t ) to share the same limit, by Lemma A.9 and Lemma A.10.Fortunately, with the monotonicity of γ c (w t ), we can prove that L t converges to zero and even characterize the rate of L t . Lemma A.14 (Rate of L t in general model). Suppose Assumption 1 holds. If there is an s such thatL(w s ) ≤ min 1 e κ+2 2n , 1 η(4ρ 2 + 2β) ,then for every t ≥ s we have1 1 L(ws) + 3ηρ 2 (t -s) ≤ L(w t ) ≤ 2 (t -s)ηγ c (w s ) 2 . That is, L(w t ) = Θ(1/t) → 0 as t → ∞.Proof of Lemma A.14. By Lemma A.3 and (12) in the proof of Lemma A.13, we know L t is decreasing andEQUATION)We will establish an upper bound for ρ t first. Note that γ c (w t ) is increasing for t ≥ s by Lemma A.13 and γ b (w t ) ≥ γ c (w t ) by Lemma A.10. By (9), we haveρ t = -log(2ne κ L t ) γ b (w t ) ≤ -log(2ne κ L t ) γ c (w t ) ≤ -log(2ne κ L t ) γ c (w s ) .Combining this with Lemma A.7, we havev t ρ t ≥ -L t log(2ne κ L t ) -log(2ne κ Lt) γ c (ws) = L t γ c (w s ).Plugging this into ( 16), we haveL t+1 -L t ≤ - η 2 L 2 t γ c (w s ) 2 , which implies ηγ(w s ) 2 2 ≤ L t -L t+1 L 2 t ≤ L t -L t+1 L t L t+1 Since L t+1 ≤ L t = 1 L t+1 - 1 L t , t ≥ s.Telescoping the sum from s to t, we have(t -s) ηγ c (w s ) 2 2 ≤ 1 L t - 1 L s ≤ 1 L t .Therefore we haveL t ≤ 2 (t -s)ηγ c (w s ) 2 .Next we show the lower bound on the risk. By Lemma A.3 we haveL t+1 -L t ≥ -η(1 + β ηL t )∥∇L t ∥ 2 ≥ -3 2 η∥∇L t ∥ 2 .Observe that under Assumption 1A,∥∇L t ∥ = 1 n n i=1 ℓ ′ (q i (t))y i ∇f (w t ; x i ) ≤ ρL t .Then we haveL t+1 -L t ≥ -η 3 2 ρ 2 L 2 t , t ≥ s. Let Lt := 3ηρ 2 2 L t , we have Ls ≤ 3ηρ 2 2 1 η(4ρ 2 +2β) ≤ 3 8 ≤ 1 2 . Furthermore, since L t decreases monotonically, Lt ≤ Ls ≤ 1 2 . The inequality becomes Lt+1 -Lt ≥ -L2 t . Therefore, let c = 1Ls and apply Lemma G.1, we have for any t ≥ s,Lt ≥ 1 c + 2(t -s) .This is equivalent toL t ≥ 1 1 Ls + 3ηρ 2 (t -s).We have completed the proof.Furthermore, we can characterize the order of ρ t in the stable phase.Lemma A.15 (Order of ρ t in general model). Suppose Assumption 1 holds. If there is s such thatL(w s ) ≤ min 1 e κ+2 2n , 1 η(4ρ 2 + 2β) ,then for t ≥ s we have ρ t = Θ(log(t)).Proof of Lemma A.15. Note that γ c (w t ) is increasing for t ≥ s by Lemma A.13 and γ b (w t ) ≥ γ c (w t ) by Lemma A.10. Therefore,ρ t ≤ -log(2ne κ L t ) γ b (w t ) ≤ -log(2ne κ L t ) γ c (w t ) ≤ -log(2ne κ L t ) γ c (w s ) .Combining this with Lemma A.14, we haveρ t ≤ log 1/L(ws)+3ηρ 2 (t-s) 2ne κ γ c (w s ) = O(log(ηt)).Besides, we have q min ≥ ι(log 1 L -log n) by Lemma A.9 and q min ≤ B 0 ρ t by Lemma A.12. Therefore we haveρ t ≥ ι(log 1 nLt ) B 0 ≥ log 1 nLt 2B 0 ≥ log (t-s)ηγ c (ws) 2 2n 2B 0 = Ω(log(t)),where the second inequality is because for ι(x) defined in (6), ι(x) ≥ x 2 for x ≥ 0.6, and the third inequality is by Lemma A.14. Combining them, we getρ t = Θ(log(t)).This completes the proof.Proof of Theorem 2.2. We prove the items one by one.• The monotonicity of L t comes from the result of Lemma A.3 directly.• Item 1 is due to Lemma A.14 .• For item 2, the monotonicity of ρ t 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 t → 0 by Lemma A.14. Then, by Lemma A.13 and Lemma A.12, we know that γ c (w t ) converges. Combining these with Lemma A.9 and Lemma A.10, we know that γ c (w t ) is an 1 + O 1/(log 1 L(wt) -multiplicative approximation ofγ(w t ), and γ(w t ) shares the same limit as γ c (w t ).In this section, we focus on the linearly separable case, that is, we work under Assumption 3. We mainly follow the idea of [Wu et al., 2024] for the proof. In detail, we consider a comparatoru := u 1 + u 2 ,where whereEQUATIONandEQUATIONConsider the following decomposition,∥w t+1 -u∥ 2 = ∥w t -u∥ 2 + 2mη⟨∇L(w t ), u -w t ⟩ + m 2 η 2 ∥∇L(w t )∥ 2 = ∥w t -u∥ 2 + 2mη ⟨∇L(w t ), u 1 -w t ⟩ =:I1(wt) +mη 2⟨∇L(w t ), u 2 ⟩ + mη∥∇L(w t )∥ 2 =:I2(wt).We aim to prove I 1 (w t ) ≤ 1 T -L(w t ) and I 2 (w t ) ≤ 0. Then we can get a bound for the average loss by telescope summing the decomposition. Here we also introduced the following vector w * :w * :=    a 1 w * . . . a n w *   We can observe that u 1 = log(γ 2 ηT )+κ αγ w * and u 2 = η 2γ w * . Lemma B.1 (A bound on I 1 (w) in the EoS phase). For u 1 defined in (17), we haveI 1 (w) := ⟨∇L(w), u 1 -w⟩ ≤ 1 γ 2 ηT -L(w).Proof. Since L is averaged over the individual losses incurred at the data (x i , y i ) n i=1 and gradient is a linear operator, it suffices to prove the claim assuming there is only a single data point (x, y). Then by Assumption 3, we have ⟨yx, w * ⟩ ≥ γ > 0. Then the loss becomesL(w) = ℓ(yf (w; x)) = ℓ y 1 m m j=1 a j ϕ(x ⊤ w (j)) .Now we expand I 1 (w):EQUATION)By definition we have J 2 = yf (w; x). As for J 1 , using ϕ ′ ≥ α and a k yx ⊤ u (k)1 ≥ 0 by Assumption 3, we haveEQUATIONPlugging in J 2 = yf (w; x) and ( 20) into (19), we getI 1 (w) = ⟨∇L(w), u 1 -w⟩ = ℓ ′ yf (w; x) (J 1 -J 2 ) ≤ ℓ ′ yf (w; x) log(γ 2 ηT ) -yf (w; x) since ℓ ′ < 0 ≤ ℓ(log(γ 2 ηT )) -ℓ(yf (w; x)) since ℓ is convex ≤ 1 γ 2 ηT -L(w).where in the last inequality, we use ℓ(x) ≤ exp(-x) and we only consider a single data point. This completes the proof.Lemma B.2 (A bound on I 2 (w) in EoS). For u 2 defined in (18), for every w,I 2 (w) := 2⟨∇L(w), u 2 ⟩ + mη∥∇L(w)∥ 2 ≤ 0.Proof. For simplicity, we defineg i (w (j) ) := ℓ ′ (y i f (w; x i ))ϕ ′ (x ⊤ i w (j) ). Note that -1 ≤ ℓ ′ (•) ≤ 0 and 0 < α ≤ ϕ ′ (•) ≤ 1, we have -1 ≤ g i (w (j) ) ≤ 0.Under this notation, we have∂L(w) ∂w i = 1 n n i=1 ℓ ′ y i f (w; x i ) y i a j m -1 ϕ ′ x ⊤ i w (j) x i = 1 n n i=1 g i w (j) a j m -1 y i x i .So we haveI 2 (w) := 2⟨∇L(w), u 2 ⟩ + mη∥∇L(w)∥ 2 = 1 m m j=1 2 n n i=1 g i (w (j) )a j y i • x ⊤ i u (j) 2 + η 1 n n i=1 g i (w (j) )a j y i x i 2 .For the term inside the bracket, we have2 n n i=1 g i (w (j) )a j y i • x ⊤ i u (j) 2 + η 1 n n i=1 g i (w (j) )a j y i x i 2 = 2 n n i=1 g i (w (j) )a j y i • x ⊤ i η 2γ a j w * + η 1 n n i=1 g i (w (j) )a j y i x i 2 since u (j) 2 := ηa j 2γ w * by (18) ≤ 2 n n i=1 g i (w (j) )a 2 j η 2γ γ + η 1 n n i=1 g i (w (j) )a j y i x i 2 since g i (•) ≤ 0 and y i x ⊤ i w * ≥ γ =η 1 n n i=1 g i (w (j) ) + 1 n n i=1 g i (w (j) )y i x i 2 since a 2 j = 1 ≤η 1 n n i=1 g i (w (j) ) + 1 n n i=1 g 2 i (w (j) ) since |g i (•)| ≤ 1 and ∥yx∥ ≤ 1 ≤0. since -1 ≤ g i (•) ≤ 0Hence, we prove that I 2 (w) ≤ 0.Theorem B.3 (A split optimization bound). For every η > 0 and u = u 1 + u 2 such thatEQUATIONu 2 :=    u (1) 2 . . . u (m) 2    , with u (j) 2 := a j η 2γ • w * , j = 1, . . . , m.we have:∥w T -u∥ 2 2mηT + 1 T T -1 k=0 L w (k) ≤ 1 + 8 log 2 (γ 2 ηT )/α 2 + 8κ 2 /α 2 + η 2 γ 2 ηT + ∥w 0 ∥ 2 mηT ,for all T .Proof. By Lemma B.1 and Lemma B.2, we have∥w t+1 -u∥ 2 = ∥w t -u∥ 2 + 2mηI 1 (w t ) + ηmI 2 (w t ) ≤ ∥w t -u∥ 2 + 2mηI 1 (w t ) ≤ ∥w t -u∥ 2 + 2mη 1 γ 2 ηT -L(w t ) .Telescoping the sum, we get∥w T -u∥ 2 2mη + T -1 t=0 L(w t ) ≤ 1 + ∥w 0 -u∥ 2 2mη .By ( 17)and ( 18), we have∥w 0 -u∥ 2 ≤ 2∥w 0 ∥ 2 + 2∥u∥ 2 ≤ 2∥w 0 ∥ 2 2 + 4∥u 1 ∥ 2 + 4∥u 2 ∥ 2 = 2∥w 0 ∥ 2 2 + 8m log(γ 2 ηT ) 2 + 8mκ 2 α 2 γ 2 + mη 2 γ 2 , which implies that ∥w T -u∥ 2 2mηt + 1 T T -1 k=0 L w (k) ≤ 1 + 8 log 2 (γ 2 ηT )/α 2 + 8κ 2 /α 2 + η 2 γ 2 ηT + ∥w 0 ∥ 2 mηT .We complete the proof.Proof of Theorem 3.2. By Theorem B.3, we have1 T T -1 k=0 L w (k) ≤ 1 + 8 log 2 (γ 2 ηT )/α 2 + 8κ 2 /α 2 + η 2 γ 2 ηT + ∥w 0 ∥ 2 mηT .This completes the proof.Lemma C.2 (A bound of ∥w t ∥). For every η, we have∥w t ∥ ≤ √ m • 2 + 8 log(γ 2 ηt)/α + 8κ/α + 4η γ + 2∥w 0 ∥.Proof of Lemma C.2. By Theorem B.3, we have∥w t -u∥ 2 2mηt ≤ ∥w t -u∥ 2 2mηt + 1 t t-1 k=0 L w (k) ≤ 1 + 8 log 2 (γ 2 ηt)/α 2 + 8κ 2 /α 2 + η 2 γ 2 ηt + ∥w 0 ∥ 2 mηt .Besides, we know that∥u∥ 2 ≤ 2∥u 1 ∥ 2 + 2∥u 2 ∥ 2 = 4m log(γ 2 ηt) 2 + 4mκ 2 α 2 γ 2 + mη 2 2γ 2 .Combining them, we have∥w t ∥ 2 ≤ 2∥w t -u∥ 2 + 2∥u∥ 2 ≤ m • 2 + 24 log 2 (γ 2 ηt)/α 2 + 24κ 2 /α 2 + 3η 2 γ 2 + 2∥w 0 ∥ 2 .Hence, we can get a bound for ∥w t ∥.∥w t ∥ ≤ √ m • 2 + 8 log(γ 2 ηt)/α + 8κ/α + 4η γ + 2∥w 0 ∥.Lemma C.3 (Gradient potential bound in the EoS phase). For every η, we have1 t t-1 k=0 G(w (k) ) ≤ ⟨w t , w * ⟩ -⟨w 0 , w * ⟩ mαγηt ≤ √ m∥w t ∥ -⟨w 0 , w * ⟩ mαγηt , t ≥ 1.Additionally, we have1 t t-1 k=0 G(w (k) ) ≤ 2 + 8 log γ 2 ηt /α + 8κ/α + 4η αγ 2 ηt + 3∥w 0 ∥ αγηt , t ≥ 1.Proof. This is from the perceptron argument [Novikoff, 1962] . Specifically,⟨w t+1 , w * ⟩ = ⟨w t , w * ⟩ -mη⟨∇L(w t ), w * ⟩ = ⟨w t , w * ⟩ -η n i=1 m k=1 a 2 k ℓ ′ (y i f (w t ; x i ))y i ϕ(x ⊤ i w (k) t )⟨x i , w * ⟩ ≥ ⟨w t , w * ⟩ -η n i=1 m k=1 a 2 k ℓ ′ (y i f (w t ; x i ))αγ ≥ ⟨w t , w * ⟩ + mαγηG(w t ).Telescoping the sum, we have1 t t-1 k=0 G(w (k) ) ≤ ⟨w t , w * ⟩ -⟨w 0 , w * ⟩ mαγηt ≤ √ m∥w t ∥ -⟨w 0 , w * ⟩ mαγηt ≤ 2 + 8 log γ 2 ηt /α + 8κ/α + 4η αγ 2 ηt + 3∥w 0 ∥ √ mαγηt . by Lemma C.2We have completed the proof.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 s such thatL(w s ) ≤ min 1 η(4 + 2 β) , 1 2e κ+2 n .Then for every t ≥ s + 1, we haveL(w t ) ≤ 2 (t -s)α 2 γ 2 . Proof. By Lemma A.3 and f (x) is 1 √ m Lipschitz and β m smooth, we have L k+1 ≤ L k -mη(1 -(2 + β)ηL(w k ))∥∇L t ∥ 2 .By Lemma C.1 and L t ≤ 1 η(4+2 β) , we haveL k+1 ≤ L k - α 2 γ 2 2 L 2 k .Multiplying 1 L 2 k in both sides, we haveα 2 γ 2 2 ≤ L t -L k+1 L 2 k ≤ 1 L k+1 - 1 L k .Taking summation for k = s, . . . , t -1, we have1 L t > 1 L t - 1 L s ≥ (t -s)α 2 γ 2 2 =⇒ L t ≤ 2 (t -s)α 2 γ 2 .At last, we will use the bound for the gradient potential to get an upper bound for the phase transition time.Proof of Theorem 4.1. Applying Lemma C.3, we have1 τ τ -1 k=0 G(w (k) ) ≤ 2 + 8 log γ 2 ητ /α + 8κ/α + 4η αγ 2 ητ + 3∥w 0 ∥ √ mαγητ ≤ 2 + 8κ/α + 8 log(γ 2 τ )/α + (4 + 8/α)η αγ 2 ητ + 3∥w 0 ∥ √ mαγητ since log(η) ≤ η, Let c 1 = 4e κ+2 , c 2 = (8 + 4 β). Note that we have 2 + 8κ/α αγ 2 ητ ≤ 1 4(c 1 n + c 2 η) if γ 2 τ ≥ 4(2 + 8κ) c 2 η + c 1 n ηα 2 8 log(γ 2 τ )/α αγ 2 ητ ≤ 1 4(c 1 n + c 2 η) if γ 2 τ ≥ 128 c 2 η + c 1 n ηα 2 log c 2 η + c 1 n η , since Lemma G.5 (4 + 8/α)η αγ 2 ητ ≤ 1 4(c 1 n + c 2 η) if γ 2 τ ≥ 48 α 2 (c 2 η + c 1 n), 3∥w 0 ∥ mαγητ ≤ 1 4(c 1 n + c 2 η) if γτ ≥ 12 α (c 2 η + c 1 n) η • ∥w 0 ∥ √ mand that the two conditions are satisfied becauseγ 2 τ := 128(1 + 4κ) α 2 max c 2 η, c 1 n, e, c 2 η + c 1 n η log c 2 η + c 1 n η , (c 2 η + c 1 n)∥w 0 ∥ η √ m ≥ max 4(2 + 8κ) c 2 η + c 1 n ηα 2 , 128 c 2 η + c 1 n ηα 2 log c 2 η + c 1 n η , 48 α 2 (c 2 η + c 1 n), 12 α (c 2 η + c 1 n) η • ∥w 0 ∥ √ m .So there exits s ≤ τ such thatG(w s ) ≤ min 1 e κ+2 4n , 1 η(8 + 4 β)Then we have L(w s ) ≤ F (w s ) ≤ 2G(w s ) ≤ 1 e κ+2 2n , 1 η(4+2 β) . We complete the proof.Proof of Corollary 4.2. The main idea is to show that τ ≤ T 2 . Note that by Theorem 4.1, we haveτ = 128(1 + 4κ) α 2 max c 2 η, c 1 n, e, c 2 η + c 1 n η log c 2 η + c 1 n η , (c 2 η + c 1 n) η • ∥w 0 ∥ √ m ,in which expression c 1 = 4e κ+2 and c 2 = (8 + 4 β). We can verify that,128(1 + 4κ) α 2 c 2 η = 128(1 + 4κ) α 2 c 2 • α 2 γ 2 256(1 + 4κ)c 2 T = T 2 , 128(1 + 4κ)c 1 n α 2 ≤ T 2 .Furthermore, we have n ≤ α 2 γ 2 T 256(1+4κ)c1 . Hence,c 2 η + c 1 n η = α 2 γ 2 T 256(1+4κ) + c 1 n α 2 γ 2 T 256(1+4κ)c2 ≤ 2 • α 2 γ 2 T 256(1+4κ) α 2 γ 2 T 256(1+4κ)c2 ≤ 2c 2 .We get that:128(1 + 4κ) α 2 • c 2 η + c 1 n η log c 2 η + c 1 n η ≤ 2 128(1 + 4κ) α 2 c 2 ln(2c 2 ) ≤ 128(1 + 4κ) α 2 4c 2 2 ≤ T 2 , 128(1 + 4κ) α 2 • (c 2 η + c 1 n) η • ∥w 0 ∥ √ m ≤ 128(1 + 4κ) α 2 • 2c 2 ∥w 0 ∥ √ m ≤ T 2 .Hence, we have τ ≤ T 2 . Applying Theorem 4.1, we haveL(w T ) ≤ 2 α 2 γ 2 η(T -τ ) ≤ 4 α 2 γ 2 ηT ≤ 2048(1 + 4κ)c 2 α 4 γ 4 T 2 = O(1/T 2 ).We have completed the proof.Proof of Theorem 4.3. The main idea is to construct an upper bound of η and apply the analysis in Theorem 2.2. Note that give w 0 = 0, we havef (w 0 ; x i ) = 1 m m k=1 a k ϕ(x ⊤ i w (k) 0 ) = s a ϕ(0),where s a = m k=1 a k /m. Therefore,[∇L(w 0 )] (k) = 1 2 ℓ ′ (s a ϕ(0)) • a k m ϕ ′ (0)x 1 + 1 2 ℓ ′ (s a ϕ(0)) • a k m ϕ ′ (0)x 2 = a k m ℓ ′ (s a ϕ(0))ϕ ′ (0) x 1 + x 2 2 = a k m ℓ ′ (s a ϕ(0))ϕ ′ (0)(γ, 1 -γ 2 4 ). Let x := 1 m ℓ ′ (s a ϕ(0))ϕ ′ (0)(γ, √ 1-γ 2 4), we havew (k) 1 = 0 -η∇[L(w 0 )] (k) = -ηa k x. Therefore, f (w 1 ; x i ) = 1 m m k=1 a k ϕ(-x ⊤ i (ηa k x)).We can notice that -x ⊤ 1 x < 0 and -x ⊤ 2 x > 0, when γ ≤ 0.1. Furthermore, we havef (w 1 ; x 1 ) = 1 m a k =1 ϕ(-x ⊤ 1 (ηx)) + 1 m a k =-1 -ϕ(x ⊤ 1 (ηx)) = 1 m a k =1 [ϕ(0) -x ⊤ 1 (ηx)ϕ ′ (ϵ 1 )] + 1 m a k =-1 [-ϕ(0) -x ⊤ 1 (ηx)ϕ ′ (ϵ 2 )] = s a ϕ(0) -ηx ⊤ 1 x 1 m [ a k =1 ϕ ′ (ϵ 1 ) + a k =-1 ϕ ′ (ϵ 2 )] ≤ s a ϕ(0) -ηx ⊤ 1 xα ϕ ′ (ϵ i ) ≥ α. Note that 1 2 ℓ(s a ϕ(0) -ηx ⊤ 1 xα) ≤ 1 2 ℓ(f (w 1 ; x 1 )) ≤ L(w 1 ) ≤ L(w 0 ) = ℓ(s a ϕ(0)).We apply Lemma G.7 to getη ≤ |s a ϕ(0)| + ln 3 x ⊤ 1 xα.We use c 3 := |saϕ(0)|+ln 3x ⊤ 1 xα. Now we know η ≤ c 3 . Furthermore, notice that∥∇L(w t )∥ ≤ L t ≤ L 0 .We get that ∥w t+1 -w t ∥ ≤ ηL 0 ≤ c 3 L 0 . Hence, |f (w t+1 ; x i ) -f (w t ; x i )| ≤ c 3 L 0 .Assume that l b = min1 e κ+2 4n , 1 η(8+4 β) and L s-1 ≥ l b , L s ≤ l b .We know thatl b ≥ min 1 e κ+2 4n , 1 c 3 (8 + 4 β) =: l c .We want to show that there is an lower bound for L s . Now thatL s = 1 2 ℓ(f (w s ; x 1 )) + ℓ(f (w s ; x 2 )) .Applying Lemma G.6, we can get thatL s ≥ exp(-c 3 L 0 )L s-1 ≥ exp(-c 3 L 0 )l b .Recall that by Lemma A.14, we haveL t ≥ 1 1 Ls + 3ηρ 2 (t -s) , t ≥ s.Combine this with ρ = 1 √ m , η = ηm and L s ≥ exp(-c 3 L 0 )l b and we getL t ≥ 1 exp(c3L0) l b + 3η(t -s) , t ≥ s.Note that when t ≤ s, L t ≥ l b . We can get a lower bound for L t byL t ≥ 1 exp(c3L0) l b t + 3ηt ≥ 1 exp(c3L0) l b t + 3c 3 t ≥ 1 exp(c3L0) lc t + 3c 3 t = c 4 t ,where c 4 = 1 exp(c 3 L 0 ) lc +3c3 depends only on {a j } m j=1 , ϕ(0), κ, β and n.Combining the analysis for two cases, we observe that when b ≥ 1, the fast loss rate is O(1/T 2 ) given finite budget T . While b < 1, the rate is O(b -3 /T 2 ). In our main results, we set b = 1 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) = 1 √ m m j=1 a j ϕ(x ⊤ w (j) ).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.Here we provide additional experiments to support our theoretical results. In Figure 3 , we show the test accuracy of 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 a 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 an open question to understand each neuron's dynamics in deep networks. Proof of Example 2.1. Recall that the two-layer neural network is defined as:f (w; x) = 1 m m j=1 a j ϕ(x T w (j) ).We can verify that if ϕ(x) is β-smooth and ρ-Lipschitz with respect to x, then f (w; x) is β/m-smooth and ρ/ √ m-Lipschitz with respect to w. This is because: Figure 4 : 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 φ(x) = (x + ϕ(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.∇L(w) = Êℓ ′ (yf (w; x))y∇f (w; x),∇ 2 L(w) = Êℓ ′′ (yf (w; x))∇f (w; x) ⊗2 + ℓ ′ (yf (w; x))y∇ 2 f (w; x),and that∇f (w; x) =     . . . 1 m a j ϕ ′ (x ⊤ w (j) )x . . .     , ∇ 2 f (w; x) =     . . . 0 0 0 1 m a j ϕ ′′ (x ⊤ w (j) )xx ⊤ 0 0 0 . . .     .Now, we will focus on the parameters of each activation function.• GELU. ϕ(x) = x • erf(1 + (x/ √ 2))/2 = x • F (x). ϕ ′ (x) = F (x) + x • f (x), ϕ ′′ (x) = 2f (x) + x • f ′ (x),where F (x), f (x) are the CDF and PDF of standard normal distribution. Note that xf (x) =x √ 2π e -x 2 /2 and (xf (x)) ′ = 1 √ 2π (1 -x 2 )e -x 2 /2 . We can find the maximum of xf (x) is1 √ 2π e -1/2 . Besides, we know that F (x), f (x) ≤ 1 and x • f ′ (x) ≤ 0. Combining them, we have ρ = 1 + e -1/2 / √ 2π and β = 2. For κ, ϕ -ϕ ′ (x)x = -x • f (x). So the bound of κ is e -1/2 / √ 2π.• Softplus. ϕ(x) = log(1 + e x ). Therefore,ϕ ′ (x) = e x 1 + e x ≤ 1, ϕ ′′ (x) = e x (1 + e x ) 2 ≤ 1. Besides, (ϕ(x) -ϕ ′ (x)x) ′ = log(1 + e x ) - e x x 1 + e x ′ = - e x x (1 + e x ) 2 . So the maximum is ϕ(0) -ϕ ′ (0)0 = log 2. Besides, when x > 1, ϕ(x) ≥ x ≥ e x x 1+e x . When x → -∞, ϕ(x) -ϕ ′ (x)x → 0. Therefore, κ = log 2. • Sigmoid. ϕ(x) = 1/(1 + e -x ). Hence, ϕ ′ (x) = e -x (1 + e -x ) 2 ≤ 1, ϕ ′′ (x) = e -2x -e -x(1 + e -x ) 3 ≤ 1.As for κ, we know that|ϕ(x) -ϕ ′ (x)x| = 1 + e -x -xe -x (1 + e -x ) 2 ≤ 1 + e -x + |x|e -x (1 + e -x ) 2 .Note that |x|e -x ≤ e -2x + 1. We have|ϕ(x) ≤ 1 + e -x + e -2x + 1 (1 + e -x ) 2 ≤ 2. • Tanh. ϕ(x) = e x -e -x e x +e -x ≤ 1. Note that ϕ ′ (x) = 1 -ϕ(x) 2 ≤ 1 ϕ ′′ (x) = 2ϕ(x) 3 -ϕ(x) ≤ 2.Besides, we know that|xϕ ′ (x)| = 4|x| (e x + e -x ) 2 ≤ 4. Hence |ϕ(x) -ϕ ′ (x)| ≤ |ϕ(x)| + |xϕ ′ (x)| ≤ 5. • SiLU. Note that ϕ ′ (x) = 1 + e -x + xe -x (1 + e -x ) 2 , ϕ ′′ (x) = (2 -x)e -x (1 + e -x ) 2 + xe -2x (1 + e -x ) 3 . Because |x|e -x ≤ e -2x + 1 and |x|e -2x ≤ e -3x + 1. We get |ϕ ′ (x)| ≤ 2 and ϕ ′′ (x)| ≤ 4. At last, |ϕ(x) -xϕ ′ (x)| = |x|e -x (1 + e -x ) 2 ≤ 1. • Huberized ReLU. It's obvious that ϕ ′ (x) ≤ 1 and β = 1/h. Note that ϕ is not second-order differentiable. At last, |ϕ(x) -xϕ ′ (x)| =    0 x < 0, x 2 /2h 0 ≤ x ≤ h, h/2 x > h.Hence, it's upper bounded by h/2. Lemma G.1. If 1 2 ≥ L 1 ≥ 1 c and L 2 ≥ L 1 -L 2 1 , we have L 2 ≥ 1 c + 2 . Proof. For function g(x) = x -x 2 , g ′ (x) = 1 -2x. If x ≤ 1 2 , then g(x) is increasing. Then g(L 1 ) ≥ g( 1 c ) = c -1 c 2 = c 2 + c -2 c 2 (c + 2) ≥ 1 c + 2 . Lemma G.EQUATIONNow for any ∥x∥ > r, let y = rx ∥x∥ and consider g(s) = f (sy) s . Then we haveg ′ (s) = ⟨∇f (sy), sy⟩ -f (sy) s 2 . Therefore, -κ s 2 ≤ g ′ (s) ≤ κ s 2 . Let s = ∥x∥/r, f (x)r ∥x∥ = g(s) = g(1) + s 1 g ′ (t)dt ≤ g(1) + s 1 κ t 2 dt ≤ g(1) + κ ≤ rC r + κ. Therefore, f (x) ≤ (C r + κ r ) • ∥x∥. Similarly, we can show that -f (x) ≤ (C r + κ r ) • ∥x∥. Therefore, for any ∥x∥ ≥ r, |f (x)| ≤ (C r + κ r ) • ∥x∥. Let D = max ∥x∥≤r |f (x)|, we have for any x, |f (x)| ≤ (C r + κ r ) • ∥x∥ + D.We have completed the proof.Lemma G.3. Fixing c > 1, then for every 0 < x ≤ 1 c , we havex ≤ -1 log(cx) .Proof of Lemma G.3. This is equivalent to show thatx log(cx) ≥ -1.Let s(x) = x log(cx), then s ′ (x) = 1 + log(cx). Hence s(x) is decreasing when 0 < x < 1 ce and is increasing when x ≥ 1 ce . The minimum of s(x) is achieved at x = 1 ce , which iss(1/(ce)) = - 1 ce ≥ -1.This completes the proof.Lemma G.4. Given 0 < b ≤ 1 2 and a > 0, we have1 + a 1 -b ≤ (1 + 2a + 2b).Proof. This is equivalent to show that(1 + a) ≤ (1 + 2a + 2b)(1 -b) = 1 + 2a + b -2ab -2b 2 .This is equivalent to 2b(a + b) ≤ (a + b).Since a + b > 0 and b ≤ 1 2 , this is true. Lemma G.5. Given c > e, we have for any x > 2c log c,log x x ≤ 1 c . Proof. It's equivalent to show that x -c log x ≥ 0. Let g(x) = x -c log x. g ′ (x) = 1 -c/x. When x > 2c log c > 2c, g ′ (x) < 0. Hence, the minimal is g(2c log c). Note that g(2c log c) = 2c log c -c log c -c log 2 -c log log c = c log c -c log 2 -c log log c = c log c 2 log c .Now we want to show that c > 2 log c. Let h(y) = y -2 log y. h ′ (y) = 1 -2/y > 0 when y > e. h(e) = e -2 > 0. Hence h(c) > h(e) > 0 and g(2c log c) > 0. This leads to g(x) > 0. Then, we complete the proof.Lemma G.6. Given ℓ(x) = log(1 + e -x ) and c > 0, we have for any x, ℓ(x + c) ≥ exp(-c)ℓ(x).Proof. Let g(x) = ℓ(x + c) -exp(-c)ℓ(x). Then, we haveg ′ (x) = -1 1 + exp(x + c) + 1 exp(c) + exp(x + c) < 0.Therefore, g(x) is monotonically decreasing. When x → ∞, we havelim x→∞ g(x) = lim x→∞ [ℓ(x + c) -exp(-c)ℓ(x)] = exp(-x -c) -exp(-c) exp(-x) = 0.Therefore, g(x) ≥ 0 for any x. Now, we complete the proof.Lemma G.7. Assume ℓ(x) = log(1 + e -x ). If ℓ(x + c) ≤ 2ℓ(x), we havec ≤ ln 3 + |x|. Proof. Note that ℓ(x + c) -2ℓ(x) = log 1 + e x+c 1 + 2e x + e 2x ≤ 0. Then, 1 + e x+c 1 + 2e x + e 2x ≤ 1 =⇒ e c ≤ 2 + e x ≤ 2 + e |x| ≤ 3e |x| . Therefore, c ≤ ln 3 + |x|.• 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. Justification: This is a theoretical paper, and we do not discuss societal impacts. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).+ log(2ne κ x) ≤ 0, log(2ne κ x) < 0, and -log 2 (2ne κ x) -log(2ne κ x) < 0. Plugging these into the previous equality, we have Φ ′′ (x) ≥ 0 when 0 < x < 1 2ne 2+κ . 


Decision: 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.

Reviews: review: paper_summary: paper_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:

- 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.
- 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.
- For linearly separable datasets, GD with an appropriately chosen step size achieves an accelerated convergence rate of $O(1/T^2)$.
strength_weakness: strength_weakness: - Despite studying two-layer neural networks, the main theorems require linear separability of datasets, except for Theorem 2.2.
- 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.

[D. Barrett and B. Dherin] IMPLICIT GRADIENT REGULARIZATION. ICLR, 2021

[Q. Li, C. Tai, and W. E] Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations. JMLR, 2019.

[M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] SGD with Large Step Sizes Learns Sparse Features. ICML, 2023.

[Y. Ren, C. Ma, and L. Ying] Understanding the Generalization Benefits of Late Learning Rate Decay. AISTATS, 2024.

[F. Chen, Z. Ren, and S. Wang] Uniform-in-time propagation of chaos for mean field Langevin dynamic. 2022.

[T. Suzuki, D. Wu, and A. Nitanda] Convergence of mean-field Langevin dynamics: time-space discretization, stochastic gradient, and variance reduction. NeurIPS, 2023.
questions: 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:

[M. Telgarsky] Feature selection and low test error in shallow low-rotation ReLU networks. ICLR, 2023.

[T. Suzuki, D. Wu, K. Oko, and A. Nitanda] Feature learning via mean-field Langevin dynamics: classifying sparse parities and beyond. NeurIPS, 2023.
- 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?
- Can Assumption 1-C be relaxed to $\kappa \leq 1$?
limitations: limitations: The limitations have been well addressed.
paper_summary: paper_summary: This paper studies GD for nearly-1-homogeneous neural networks with large stepsizes. It provides two main results.
 

The 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:
     - the training loss converges to 0 at a $\frac{1}{t}$ rate
     - the normalized margin "increases", and hence converges

The 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}$).
strength_weakness: strength_weakness: 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. 
Actually, 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. 

Additionally, 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. 
Actually, 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.
questions: questions: - Is there a specific mathematical challenge (eg wrt Lyu and Li) to prove a KKT point convergence type of result?
- Same question for extending $1$-homogeneity to $L$-homogeneity for any positive integer $L$
- Can we say more about the normalized margin in the restricted setting of Theorem 4.1 ?
limitations: limitations: N/A
paper_summary: paper_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
1) normalized margin grows nearly monotonically in the second phase
2) in the oscillatory EOS phase the average empirical risk decreases under certain conditions
3) With a larger step-size, the GD undergoing phase transition is more efficient.
strength_weakness: strength_weakness: 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. 
2) 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.
questions: 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? 
2) 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.
3) 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: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: The role and effect of large learning rates, especially compared to small learning rates, may need more illustration.
questions: questions: 1. What is the initial condition in Theorem 3.2?
2. 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: 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.

Comments: title: A technical comparison with [Wu et al., 2024]
Reviewer response-1
comment: I thank the authors for their response. I updated my score.
Thank you for the response. I will maintain my positive evaluation.
I would like to thank the authors for the rebuttal. Please consider adding the discussions on the initial condition in Theorem 3.2 to the paper. I will maintain my score.
Thank you for the suggestions. We make a detailed  technical comparison with [Wu et al., 2024] below. We will add these discussions in the revision. 


Our stable phase analysis shows the margin improvement in non-homogenous networks. In comparison, the stable phase analysis in [Wu et al., 2024] only concerns the convergence of the loss. As margin improvement is harder to show (especially in non-homogenous cases), the techniques of [Wu et al., 2024] are insufficient to achieve our goal. To achieve our goal, we analyze the evolvements of several modified versions of the margin. None of these quantities appear in [Wu et al., 2024]. So our techniques here are significantly different from [Wu et al., 2024]. 

Regarding our EoS and acceleration analysis, we use tools from [Wu et al., 2024], as there are not many tools that can deal with large stepsizes. Besides extending their results from linear models to networks, we make two innovations in our analysis. First, our comparator $u_1$ (see Equation 17 in Appendix B) contains an extra component to accommodate the non-homogeneity of the predictor. Second, our Lemma C.4 uses a sharper convergence analysis, which removes some logarithmic factors in Theorem 4.1 and Corollary 4.2, comparing to these results in [Wu et al., 2024]. For instance, our Corollary 4.2 gets $O(1/T\^2)$ while their Corollary 2 gets $O(\\log(T)\^2 / T\^2)$. This is mentioned in Lines 224-227, but we will emphasize it more in the revision.
I thank the authors for their response. Since most of the comments seem addresses, I think Q-2 could have been addressed with more details rather than "analysis for nonlinear networks requires nontrivial efforts" and "use techniques significantly different from those of [Wu et al., 2024]". I think this question is important since the nature of results overlap with those of https://arxiv.org/pdf/2402.15926. I would suggest that the authors explain or mention how different are the proof techniques from https://arxiv.org/pdf/2402.15926  atleast briefly in the manuscript.
I thank the authors for their answer. I now understand more clearly how these possible extensions still require extensive work. I think the additional discussions in the revised version will improve the quality of the paper.

I raise my score in consequence