Title: Large Stepsize Gradient Descent for Non-Homogeneous Two-Layer Networks: Margin Improvement and Fast Optimization
Venue: NeurIPS 2024
Accepted: yes
Keywords: Deep Learning Theory

Large stepsize

Optimization

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

Large Stepsize Gradient Descent for Non-Homogeneous Two-Layer Networks: Margin Improvement and Fast Optimization

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.

1 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.

(a) Empirical risk

(b) Normalized margin

(c) Test accuracy

101

0.00020

0.92

=0.2

= 0.2

100

=0.005

0.91

=0.005

0.00015

101

margin

0.90

SSOl

Test accuracy

Training

102

Normalized

0.00010

0.89

103

0.00005

0.88

10-4

0.87

105

= 0.2

0.00000

0.86

10-6

=0.005

100

101

102

103

104

105

-0.00005

0

5

10

15

20

0.85

0.0

2.5

5.0

7.5 10.012.5 15.017.5

20.0

Iterations

Iterations

1e4

Iterations

1e4

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 ˜ O (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 as

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,

We are interested in a nonlinear predictor f and a large stepsize ˜ η . A notable example in our theory is two-layer networks with Lipschitz, smooth, and nearly homogenous activations (see (2)). Note that minimizing L ( w ) is a non-convex problem in general.

Observation. Empirically, large stepsize GD often undergoes a phase transition, where the empirical risk defined in (1) oscillates in the first phase but decreases monotonically in the second phase (see empirical evidence in Appendix A in [Cohen et al., 2020] and a formal proof in [Wu et al., 2024] for linear predictors). This is illustrated in Figure 1. We follow Wu et al. [2024] and call the two phases the edge of stability (EoS) phase [name coined by Cohen et al., 2020] and the stable phase, respectively.

Contributions. We prove the following results for large stepsize GD for training nonlinear predictors under logistic loss.

For Lipschitz and smooth predictor f trained by GD with stepsize ˜ η , we show that as long as the empirical risk is below a threshold depending on ˜ η , GD monotonically decreases the empirical risk (see Theorem 2.2). This result extends the stable phase result in Wu et al. [2024] from linear predictors to nonlinear predictors, demonstrating the generality of the existence of a stable phase.

Assuming that GD enters the stable phase, if in addition the preditor has a bounded homogenous error (see Assumption 1C), we show that the normalized margin induced by GD, min i y i f ( w t ; x i ) / ∥ w t ∥ , nearly monotonically increases (see Theorem 2.2). To the best of our knowledge, this is the first characterization of implicit bias of GD for non-homogenous predictors. In particular, our theory covers two-layer networks with commonly used activations functions (which are often non-homogenous) that cannot be covered by existing results [Lyu and Li, 2020, Ji and Telgarsky, 2020, Chizat and Bach, 2020].

Under additional technical assumptions (the dataset is linearly separable and the derivative of the activation function is bounded away from zero), we show that the initial EoS phase must stop in O (˜ η ) steps and GD transits to the stable phase afterwards. Furthermore, by choosing a suitably large stepsize, GD achieves a ˜ O (1 /t 2 ) empirical risk after t steps. In comparison, GD that converges monotonically incurs an Ω(1 /t ) risk. This result indicates an optimization benefit of using large stepsize and generalizes the results in [Wu et al., 2024] from linear predictors to neural networks.

2 Stable Phase and Margin Improvement

In this section, we present our results for the stable phase of large stepsize GD in training nonlinear predictors. Specifically, our results apply to nonlinear predictors that are Lipschitz, smooth, and nearly homogeneous, as described by the following assumption.

Assumption 1 (Model conditions) . Consider a predictor f ( w ; x i ) , where x i is one of the feature vectors in the training set.

Lipschitzness . Assume there exists ρ > 0 such that for every w , sup i ∥∇ w f ( w ; x i ) ∥ ≤ ρ .

Smoothness . Assume there exists β > 0 such that for all w , v ,

Near homogeneity . Assume there exists κ > 0 such that for every w ,

Assumptions 1A and 1B are commonly used conditions in the optimization literature. Note that Assumption 1B implies continuous differentiability, thus ruling out networks with ReLU activation function. The continuous differentiability is only used in our current stable phase analysis. We conjecture it can be relaxed using subdifferentiability [Lyu and Li, 2020] for allowing ReLU networks.

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

where we assume a j ∈ {± 1 } are fixed and w ∈ R md , the stack of ( w ( j ) ) j m =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 ,

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) as

A 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 that

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 that

we 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

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 have

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 min i 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.

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,

Derivative condition . Assume there exists 0 < α < 1 such that α ≤ | ϕ ′ ( z ) | ≤ 1 .

Smoothness . Assume there exists ˜ β > 0 such that for all x, y ∈ R , | ϕ ′ ( x ) -ϕ ′ ( y ) | ≤ ˜ β | x -y | .

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 / 4 h .

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 have

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.

4 Phase Transition and Fast Optimization

For two-layer networks trained by large stepsize GD, Theorem 3.2 shows that the average risk must decrease over time. Combining this with Theorem 2.2, GD must enter the stable phase in finite steps, and the loss must converge while the normalized margin must improve.

However, 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

where c 1 := 4 e κ +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 and

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 that

where c 1 := 4 e κ +2 and c 2 := (8 + 4 ˜ β ) are as in Theorem 4.1. Then for stepsize ˜ η := ηm , where

Theorem 4.1 and Corollary 4.2 extend Theorem 1 and Corollary 2 in Wu et al. [2024] from linear predictors to two-layer networks. Another notable difference is that we obtain a sharper stable phase bound (and thus a better acceleration bound) compared to theirs, where we remove a logarithmic factor through a more careful analysis.

Corollary 4.2 suggests an accelerated risk bound of O (1 /T 2 ) by choosing a large stepsize that balances EosS and stable phases. We also show the following lower bound, showing that such acceleration is impossible if (GD) does not enter the EoS phase. The proof is included in Appendix C.3.

Theorem 4.3 (Lower bound in the classical regime) . Consider (GD) with initialization w 0 = 0 and stepsize ˜ η > 0 for a two-layer network (2) satisfying Assumption 2. Suppose the training set is given by

It is clear that ( x i , y i ) i =1 , 2 satisfy Assumption 3. If ( L ( w t )) t ≥ 0 is non-increasing, then

where c 0 > 0 is a function of ( α, ϕ, x 1 , x 2 , γ, κ, β ) but is independent of t and ˜ η .

we have τ ≤ T/ 2 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 following

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.

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 (NTK) scaling b = √ m [Du et al., 2018, Jacot et al., 2018] give the same acceleration effect. An NTK analysis of large stepsize is included in [Wu et al., 2024] and their conclusion is consistent with ours. Finally, we remark that our analysis holds for any width m and uses techniques different from the mean-field or NTK methods. However, our acceleration analysis only allows linearly separable datasets.

5 Experiments

We conduct three sets of experiments to validate 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:

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.

Margin of individual neurons. It is important to note that while the normalized margin behaves as expected, the margin for individual neurons may not increase and can remain negative, even when the dataset is linearly separable. A detailed example illustrating this is provided in Appendix E.

6 Related Works

In this section, we discuss related papers.

(a) Empirical risk, XOR.

(b) Asymptotic rate, XOR.

(c) Normalized margin, XOR.

(d) Empirical risk, CIFAR-10.

(e) Asymptotic rate, CIFAR-10.

(f) Normalized margin, CIFAR-10.

0.130

= 64

= 64

100

8=↓

=32

=16

103

= 32

=16

=8

0.125

0.120

10-1

-(²M)7

0.115

(2M)7

102

102

0.110

0.105

= 64

10-3,

0.100

= 32

=16

10-4

101

0.095

=8

100

101

102

103

104

0

2000

4000

6000

8000

10000

0.090

0

2000

4000

6000

8000

10000

t

t

t

0.02

=2

=2

101

=1

=1

= 0.5

103

= 0.5

0.01

100

= 0.25

= 0.25

2↓- (M)7

0.00

(2M)7

10-1

102

0.01

10-2

=2

101

=1

0.02

= 0.5

10-3

= 0.25

0.03

100

101

102

103

104

0

2000

4000

6000

8000

10000

0

2000

4000

6000

8000

10000

t

t

t

Small stepsize and implicit bias. For logistic regression on linearly separable data, Soudry et al. [2018], Ji and Telgarsky [2018] showed that the direction of small stepsize GD converges to the max-margin solution. Their results were later extended by Gunasekar et al. [2017, 2018], Nacson et al. [2019c,a,b], Ji et al. [2021], Lyu and Li [2020], Ji and Telgarsky [2020], Chizat and Bach [2020], Chatterji et al. [2021], Kunin et al. [2022] to other algorithms and non-linear models. However, in all of their analysis, the stepsize of GD needs to be small such that the empirical risk decreases monotonically. In contrast, our focus is GD with 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, 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 .

Comparison with 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, 2024]. The works by Wu et al. [2023, 2024] directly motivate our paper. In particular, for logistic regression on linearly separable data, Wu et al. [2023] showed margin maximization of GD with large stepsize and Wu et al. [2024] showed fast optimization of GD with large stepsize. Our work can be viewed as an extension of [Wu et al., 2023, 2024] from linear predictors 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.

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

References

Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. In International Conference on Machine Learning , pages 247-257. PMLR, 2022.

Kwangjun Ahn, Sebastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via edge of stability. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Maksym Andriushchenko, Aditya Vardhan Varre, Loucas Pillaud-Vivien, and Nicolas Flammarion. Sgd with large step sizes learns sparse features. In International Conference on Machine Learning , pages 903-925. PMLR, 2023.

David GT Barrett and Benoit Dherin. Implicit gradient regularization. arXiv preprint arXiv:2009.11162 , 2020.

Alon Brutzkus, Amir Globerson, Eran Malach, and Shai Shalev-Shwartz. Sgd learns overparameterized networks that provably generalize on linearly separable data. In International Conference on Learning Representations , 2018.

Niladri S Chatterji, Philip M Long, and Peter Bartlett. When does gradient descent with logistic loss interpolate using deep networks with smoothed relu activations? In Conference on Learning Theory , pages 927-1027. PMLR, 2021.

Fan Chen, Zhenjie Ren, and Songbo Wang. Uniform-in-time propagation of chaos for mean field langevin dynamics. arXiv preprint arXiv:2212.03050 , 2022.

Lei Chen and Joan Bruna. Beyond the edge of stability via two-step gradient updates. In International Conference on Machine Learning , pages 4330-4391. PMLR, 2023.

Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on learning theory , pages 1305-1338. PMLR, 2020.

Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations , 2020.

Alex Damian, Eshaan Nichani, and Jason D Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In The Eleventh International Conference on Learning Representations , 2022.

Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations , 2018.

Spencer Frei, Yuan Cao, and Quanquan Gu. Provable generalization of sgd-trained neural networks of any width in the presence of adversarial label noise. In International Conference on Machine Learning , pages 3427-3438. PMLR, 2021.

Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in 1 hour. arXiv preprint arXiv:1706.02677 , 2017.

Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems , 30, 2017.

Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in Neural Information Processing Systems , 31, 2018.

Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: Closing the generalization gap in large batch training of neural networks. Advances in Neural Information Processing Systems , 30, 2017.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems , 31, 2018.

Ziwei Ji and Matus Telgarsky. Risk and parameter convergence of logistic regression. arXiv preprint arXiv:1803.07300 , 2018.

Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems , 33:17176-17186, 2020.

Ziwei Ji, Nathan Srebro, and Matus Telgarsky. Fast margin maximization via dual acceleration. In International Conference on Machine Learning , pages 4860-4869. PMLR, 2021.

Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations , 2016.

Lingkai Kong and Molei Tao. Stochasticity of deterministic gradient descent: Large learning rate for multiscale objective function. Advances in Neural Information Processing Systems , 33:2625-2638, 2020.

Itai Kreisler, Mor Shpigel Nacson, Daniel Soudry, and Yair Carmon. Gradient descent monotonically decreases the sharpness of gradient flow solutions in scalar networks and beyond. In Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 17684-17744. PMLR, 23-29 Jul 2023.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, 2009.

Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maximum margin bias of quasi-homogeneous neural networks. In The Eleventh International Conference on Learning Representations , 2022.

Miao Lu, Beining Wu, Xiaodong Yang, and Difan Zou. Benign oscillation of stochastic gradient descent with large learning rate. In The Twelfth International Conference on Learning Representations , 2023.

Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations , 2020.

Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. Advances in Neural Information Processing Systems , 35: 34689-34708, 2022.

Chao Ma, Daniel Kunin, Lei Wu, and Lexing Ying. Beyond the quadratic approximation: The multiscale structure of neural network loss landscapes. arXiv preprint arXiv:2204.11326 , 2022.

Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry. Lexicographic and depth-sensitive margins in homogeneous and non-homogeneous deep models. In International Conference on Machine Learning , pages 4683-4692. PMLR, 2019a.

Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 3420-3428. PMLR, 2019b.

Mor Shpigel Nacson, Nathan Srebro, and Daniel Soudry. Stochastic gradient descent on separable data: Exact convergence with a fixed learning rate. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 3051-3059. PMLR, 2019c.

Yurii Nesterov. Lectures on Convex Optimization . Springer Publishing Company, Incorporated, 2nd edition, 2018. ISBN 3319915770.

Albert BJ Novikoff. On convergence proofs on perceptrons. In Proceedings of the Symposium on the Mathematical Theory of Automata . New York, NY, 1962.

Yinuo Ren, Chao Ma, and Lexing Ying. Understanding the generalization benefits of late learning rate decay. In International Conference on Artificial Intelligence and Statistics , pages 4465-4473. PMLR, 2024.

Mei Song, Andrea Montanari, and P Nguyen. A mean field view of the landscape of two-layers neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665-E7671, 2018.

Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research , 19(1): 2822-2878, 2018.

Taiji Suzuki, Denny Wu, and Atsushi Nitanda. Convergence of mean-field langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction. arXiv preprint arXiv:2306.07221 , 2023.

Yuqing Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogeneity: Convergence and balancing effect. In International Conference on Learning Representations , 2022a.

Yuqing Wang, Zhenghao Xu, Tuo Zhao, and Molei Tao. Good regularity creates large learning rate implicit biases: edge of stability, balancing, and catapult. arXiv preprint arXiv:2310.17087 , 2023.

Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along GD trajectory: Progressive sharpening and edge of stability. Advances in Neural Information Processing Systems , 35:99839994, 2022b.

Jingfeng Wu, Vladimir Braverman, and Jason D. Lee. Implicit bias of gradient descent for logistic regression at the edge of stability. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Jingfeng Wu, Peter L Bartlett, Matus Telgarsky, and Bin Yu. Large stepsize gradient descent for logistic loss: Non-monotonicity of the loss improves optimization efficiency. Conference on Learning Theory , 2024.

Lei Wu and Chao Ma. How SGD selects the global minima in over-parameterized learning: A dynamical stability perspective. Advances in Neural Information Processing Systems , 31, 2018.

Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of-stability training dynamics with a minimalist example. In The Eleventh International Conference on Learning Representations , 2022.

A Stable Phase Analysis

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:

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.

A.1 Decrease of the Loss

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:

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

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

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

Particularly, this indicates that if L ( 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

Taking an average over all data points, we have

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:

Taking the average over all data points, we have

Now we have completed the proof.

A.2 Increase of the Parameter Norm

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:

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:

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.,

Then, we have an exponential-loss-like decomposition of L t :

These two functions ψ, ι will help us to analyze the lower bound of v t . First, we list some properties of ψ and ι here.

which is equivalent to

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

For the fourth property, notice that

The denominator is positive and decreasing since

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

Proof of Lemma A.5. Recall that

Let y = e -x . We have

Define s ( y ) := ι ( x ) ι ′ ( x ) -x -log log 2 . Then,

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,

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 have

Proof of Lemma A.6. We use (7) to get

Then, 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 2 ne κ , then

Proof of Lemma A.7. By definition, we have

Applying Lemma A.6 and Lemma A.4, we have

Then we can apply Lemma A.4 to get

Putting the above two inequalities together, we have

Plugging this back to the bound of v t , we get

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 2 ne κ , 1 ˜ η (4 ρ 2 +2 β ) } , then

Proof of Lemma A.8. By definition, we have

where the last inequality is by Lemma A.7. Besides,

This completes the proof.

A.3 Convergence of the Margin

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

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]

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 as

and a modified margin as

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 2 n , we have

and

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 that

Therefore we have

Using x 2 ≤ log(1 + x ) ≤ x for 0 ≤ x ≤ 1 , we get

Then we can bound L t by

which is equivalent to

Dividing both sides by ρ t proves the second claim:

For the last claim, we only need to show that ρ t → ∞ . This is because if L t → 0 , we have for any i ∈ [ 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 2 ne κ +2 , there exists a constant c such that

Proof of Lemma A.10. To prove the first two inequalities, notice that

Using the above, (8) to (10), we have

Then, we will prove the remaining inequality. First, we have

By the definitions of

Since

q

min

(

t

)

¯

γ, γ

-

log(

b

L

t

, γ

)

c

by Lemma A.9

To simplify the notation, we let c 1 := 1 + (4 ρ 2 +2 β )˜ η and c 2 = log(2 ne κ ) . Since L t ≤ 1 2 ne κ +2 ⇒ -log(2 ne κ L t ) ≥ 2 > 1 , we have

1 + (4 ρ 2 +2 β )˜ η -log(2 ne κ L t ) = c 1 -log(2 ne κ L t ) ≤ c 1 . Besides, given x < c , we have e x ≤ 1 + e c x . Therefore, exp ( 1 + (4 ρ 2 +2 β )˜ η -log(2 ne κ L t ) ) = exp ( c 1 -log(2 ne κ L t ) ) ≤ 1 + c 1 exp( c 1 ) -log(2 ne κ L t ) . Plugging this into the bound for ¯ γ ( w t ) /γ c ( w t ) , we get ¯ γ ( w t ) γ c ( w t ) = ( 1 + c 2 -log(2 ne κ L t ) ) · exp ( c 1 -log(2 ne κ L t ) ) ≤ ( 1 + c 2 -log(2 ne κ L t ) ) · ( 1 + exp( c 1 ) c 1 -log(2 ne κ L t ) ) ≤ 1 + c 2 +exp( c 1 ) c 1 + c 2 c 1 exp( c 1 ) -log(2 ne κ L t ) Since -log(2 ne κ 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 -c 2 is decreasing when x ≥ c 2 +1 , we have

Plug this inequality into the previous bound for ¯ γ ( w t ) /γ c ( w t ) and we get

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

Lemma A.11 (Convexity of Φ ) . The function Φ( x ) defined in (10) is convex for 0 < x < 1 2 ne 2+ κ .

and that

Note that when x ≤ 1 2 ne 2+ κ , we have log(2 ne κ 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 2 ne κ , 1 ˜ η (4 ρ 2 +2 β ) } for t ≥ s , there exists B 0 such that

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 ,

Besides, by Lemma A.9, we have

Hence,

Let B 0 = C ρ s ,κ + log 2 n ρ 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

Therefore, we have then for t ≥ s , we have

L t +1 ≤ L t .

v t ≥ -L t log(2 ne κ 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.

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 that

Proof of Lemma A.13. The first claim is by Lemma A.3 and induction. The second and the third claims are consequences of Lemmas A.7 and A.8, respectively. We now prove the last claim. By Lemma A.3, we have

Multiplying both sides by 2 v t 1 -1 2 log(2 ne κ L t ) 1 -(2 ρ 2 + β )˜ ηL t > 0 , we get

Using the above we get

Recall that

For ν t , we have

Then we can decompose ∥∇ L t ∥ 2 as

Plugging this into (11) and dividing both two sides by 2 v 2 t , we have

By Lemma A.7, we have v t ≥ -L t log(2 ne κ L t ) . Define

Then we have

We are going to show that Ψ( x ) ≤ -Φ ′ ( x ) . Note that when 0 < x ≤ min { 1 e κ +2 2 n , 1 ˜ η (4 ρ 2 +2 β ) } , we have log(2 ne κ x ) < 0 and 1 -(2 ρ 2 + β )˜ ηx ) ≥ 1 2 > 0 . Therefore, we have

To get an upper bound of Ψ( x ) , we just need an upper bound of J . Let a := -1 2 log(2 ne κ x ) > 0 and b := (2 ρ 2 + β )˜ ηx ∈ (0 , 1 / 2] . Then we invoke Lemma G.4 to get

Recall that x ≤ 1 e κ +2 2 n ≤ 1 2 ne κ and 2 ne κ ≥ 1 . Then we apply Lemma G.3 to get x ≤ -1 log(2 ne κ x ) . Plugging this into the bound of J , we get

Plugging (15) into (14), we have

which verifies that Ψ( x ) ≤ -Φ ′ ( x ) . By this and (13), we have

where φ ( x ) = -log x = log(1 /x ) . Recall that for 0 < x ≤ 1 2 ne κ +2 , Φ( x ) is convex by Lemma A.11. By convexity of φ and Φ , we have

By the definition of γ c in (10), this can be rewritten as

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.

A.4 Sharp rates of Loss and Parameter Norm

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 that

then for every t ≥ s we have

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 and

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

Combining this with Lemma A.7, we have

Plugging this into (16), we have

which implies

Telescoping the sum from s to t , we have

Therefore we have

Next we show the lower bound on the risk. By Lemma A.3 we have

Observe that under Assumption 1A,

Then we have

Let ˜ L t := 3˜ ηρ 2 2 L t , we have ˜ L s ≤ 3˜ ηρ 2 2 1 ˜ η (4 ρ 2 +2 β ) ≤ 3 8 ≤ 1 2 . Furthermore, since L t decreases monotonically, ˜ L t ≤ ˜ L s ≤ 1 2 . The inequality becomes

Therefore, let c = 1 ˜ L s and apply Lemma G.1, we have for any t ≥ s ,

This is equivalent to

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 that

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,

Combining this with Lemma A.14, we have

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

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

then for t ≥ s we have

This completes the proof.

A.5 Proof of Theorem 2.2

Proof of Theorem 2.2. We prove the items one by one.

The monotonicity of L 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 ( w t ) )) -multiplicative approximation of ¯ γ ( w t ) , and ¯ γ ( w t ) shares the same limit as γ c ( w t ) .

B EoS Phase Analysis

In this section, we focus on the linearly separable case, that is, we work under Assumption 3. We mainly follow the idea of [Wu et al., 2024] for the proof. In detail, we consider a comparator

where where

and

Consider the following decomposition,

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 ∗ :

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 have

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

Now we expand I 1 ( w ) :

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 have

Then the loss becomes

Note that

-

1

ℓ

′

(

·

)

0

and

Under this notation, we have

So we have

For the term inside the bracket, we have

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 ,

Proof. For simplicity, we define

g

i

0

(

w

(

j

< α

)

) :=

ϕ

′

-

ℓ

′

(

(

y

i

f

·

1

)

(

1

g

i

w

;

x

i

))

ϕ

′

, we have

(

w

(

j

)

)

0

(

.

x

⊤

i

w

(

j

)

)

.

Hence, we prove that I 2 ( w ) ≤ 0 .

Theorem B.3 (A split optimization bound) . For every η > 0 and u = u 1 + u 2 such that

and we have:

for all T .

Proof. By Lemma B.1 and Lemma B.2, we have

Telescoping the sum, we get

By (17)and (18), we have

which implies that

We complete the proof.

B.1 Proof of Theorem 3.2

Proof of Theorem 3.2. By Theorem B.3, we have

This completes the proof.

C Phase Transition Analysis

In this section, we will analyze the phase transition. In detail, we follow the idea of [Wu et al., 2024] and apply the perceptron argument [Novikoff, 1962] to locate the phase transition time. Compare to the previous EoS phase analysis, we need an extra assumption on the smoothness of the activation function, which is the Assumption 2B.

To proceed, let us define the following quantities for the GD process:

Due to the self-boundedness of the logistic function, we can show that G ( w ) , L ( w ) , F ( w ) are equivalent in the following sense.

Lemma C.1 (Equivalence of G,L,F ) .

Proof. The first claim is by the property of the logistic loss. For the second one,

Besides, we have

For the third claim, by the assumption, we have

which implies that

Therefore,

We complete the proof.

The key ingredient of the phase transition analysis is the following lemma. The main idea is to consider the gradient potential G ( w ) instead of the loss function L ( w ) in EoS phase. And this will decrease the order of the bound of phase transition time from ˜ O ( η 2 ) to ˜ O ( η ) .

Lemma C.2 (A bound of ∥ w t ∥ ) . For every η , we have

Proof of Lemma C.2. By Theorem B.3, we have

Besides, we know that

Combining them, we have

Hence, we can get a bound for ∥ w t ∥ .

Lemma C.3 (Gradient potential bound in the EoS phase) . For every η , we have

Additionally, we have

This

Proof. This is from the perceptron argument [Novikoff, 1962]. Specifically,

Telescoping the sum, we have

We 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 that

Then for every t ≥ s +1 , we have

Proof. By Lemma A.3 and f ( x ) is 1 √ m Lipschitz and ˜ β m smooth, we have

By Lemma C.1 and L t ≤ 1 η (4+2 ˜ β ) , we have

Multiplying 1 L 2 k in both sides, we have

Taking summation for k = s, . . . , t -1 , we have

At last, we will use the bound for the gradient potential to get an upper bound for the phase transition time.

C.1 Proof of Theorem 4.1

Proof of Theorem 4.1. Applying Lemma C.3, we have

Let c 1 = 4 e κ +2 , c 2 = (8 + 4 ˜ β ) . Note that we have

and that the two conditions are satisfied because

So there exits s ≤ τ such that

Then we have L ( w s ) ≤ F ( w s ) ≤ 2 G ( w s ) ≤ { 1 e κ +2 2 n , 1 η (4+2 ˜ β ) } . We complete the proof.

C.2 Proof of Corollary 4.2

Proof of Corollary 4.2. The main idea is to show that τ ≤ T 2 . Note that by Theorem 4.1, we have

in which expression c 1 = 4 e κ +2 and c 2 = (8 + 4 ˜ β ) . We can verify that,

Furthermore, we have n ≤ α 2 γ 2 T 256(1+4 κ ) c 1 . Hence,

We get that:

Hence, we have τ ≤ T 2 . Applying Theorem 4.1, we have

We have completed the proof.

C.3 Proof of Theorem 4.3

Proof of Theorem 4.3. The main idea is to construct an upper bound of η and apply the analysis in Theorem 2.2. Note that give w 0 = 0 , we have

where s a = ∑ m k =1 a k /m . Therefore,

Let ¯ x := 1 m ℓ ′ ( s a ϕ (0)) ϕ ′ (0)( γ, √ 1 -γ 2 4 ) , we have ( k )

1

w

= 0

-

η

∇

[

L

(

w

0

)]

=

-

ηa

k

¯

x

.

Therefore,

We can notice that -x ⊤ 1 ¯ x < 0 and -x ⊤ 2 ¯ x > 0 , when γ ≤ 0 . 1 . Furthermore, we have

Note that

We know that

We use c 3 := | s a ϕ (0) | +ln3 x ⊤ 1 ¯ x α . Now we know η ≤ c 3 . Furthermore, notice that

∥∇

L

(

w

t

)

We get that ∥ w t +1 -w t ∥ ≤ ηL 0 ≤ c 3 L 0 . Hence,

|

f

(

w

t

+1

;

x

i

)

-

Assume that l b = min { 1 e κ +2 4 n , 1 η ( 8+4 ˜ β ) } and

L

s

-

1

l

b

,

L

L

t

w

t

;

x

s

i

)

L

0

.

l

b

.

We want to show that there is an lower bound for L s . Now that

Applying Lemma G.6, we can get that

Recall that by Lemma A.14, we have

Combine this with ρ = 1 √ m , ˜ η = ηm and L s ≥ exp( -c 3 L 0 ) l b and we get

Note that when t ≤ s , L t ≥ l b . We can get a lower bound for L t by

where c 4 = 1 exp( c 3 L 0 ) l c +3 c 3 depends only on { a j } j m =1 , ϕ (0) , κ, ˜ β and n .

f

(

(

k

)

c

3

L

0

.

D Scaling and Homogenous Error

In this section, we consider different scaling of two-layer networks. We add a scaling factor b into the model, i.e.,

We will show that given a limited computation budget T (total iterations), larger b and a corresponding best ˜ η = η · m will achieve the same best rate as b = 1 , i.e., O (1 /T 2 ) . While for smaller b , the rate is O ( b -3 /T 2 ) . Before we present the analysis, here are the bounds with b and ˜ η = m · η following the process of Lemma C.3:

Case when b ≥ 1 . Given the previous bounds, we have the following results following the idea in Appendix C:

Gradient potential bound: G ( w t ) ≤ C t for all t ≥ 0 ,

Phase transition threshold: G ( w s ) ≤ min { 1 / 4 e κ +2 n, 1 /η (8 ρ 2 b 2 +4 ˜ βb ) } ,

Stable phase bound: L ( w t ) ≤ 2 Cb 2 η ( t -s ) ,

where C depends on α, γ . Combine the first two arguments and assume η (8 ρ 2 b 2 +4 ˜ βb ) ≥ 4 e κ +2 n . We get s ≤ Cη (8 ρ 2 b 2 +4 ˜ βb ) . Plug this into the third bound. We have

It's obvious that the best η = T 16 ρ 2 b 2 C +8 ˜ βbC . Hence,

Then, the rate is still O (1 /T 2 ) .

Case when b < 1 . Similarly, we can get the following bounds:

Gradient potential bound: G ( w t ) ≤ Cb -2 t for all t ≥ 0 ,

Phase transition threshold: G ( w s ) ≤ min { 1 / 4 e κ +2 n, 1 /η (8 ρ 2 b 2 +4 ˜ βb ) } ,

Stable phase bound: L ( w t ) ≤ 2 Cb 2 η ( t -s ) ,

where C depends on α, γ . Without loss of generality, we can assume η (8 ρ 2 b 2 +4 ˜ βb ) ≥ 4 e κ +2 n , since η can be small enough. Then, we have

It's obvious that the best η = T 16 ρ 2 C +8 ˜ βb -1 C . Hence,

Then, we can get

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:

The model falls into the NTK regime. The loss threshold will be related to m , but the loss rate is the same as that of the mean-field scaling.

E Extra Experiments

Here 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.

0.80

0.75

2

accurac

0.70

t

0.65

=2

0.60

=1

= 0.5

0.55

= 0.25

0

2000

4000

6000

8000

10000

t

F Additional Proofs

F.1 Proof of Example 2.1

Proof of Example 2.1. Recall that the two-layer neural network is defined as:

Wecan 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:

103

102

101

(M)7

L(Wt)nt

100

10-1

10-2

0

10000

20000

30000

40000

50000

Iterations

(a) Training loss, synthetic dataset.

Normalizedmargins

0.0

-0.2

0.4

-0.6

Neuron 1

Neuron2

-0.8

Model

Maximal margin

1.0

0

10000

20000

30000

40000

50000

Iterations

(b) Normalized Margins, synthetic dataset.

and that

Now, we will focus on the parameters of each activation function.

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 ) is 1 √ 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,

Besides,

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,

As for κ , we know that

Note that | x | e -x ≤ e -2 x +1 . We have

Tanh . ϕ ( x ) = e x -e -x e x + e -x ≤ 1 . Note that

Because | x | e -x ≤ e -2 x +1 and | x | e -2 x ≤ e -3 x +1 . We get | ϕ ′ ( x ) | ≤ 2 and ϕ ′′ ( x ) | ≤ 4 . At last,

Huberized ReLU . It's obvious that ϕ ′ ( x ) ≤ 1 and β = 1 /h . Note that ϕ is not second-order differentiable. At last,

Hence, it's upper bounded by h/ 2 .

F.2 Proof of Example 3.1

Proof of Example 3.1. Because for activation functions in Example 2.1, β ≤ 4 and ρ ≤ 2 . Hence, for ˜ ϕ ( x ) = cx + (1 -c ) ϕ ( x ) / 4 , ˜ β = 1 and ρ = 1 . Besides, since 0 . 5 < c < 1 , we must have ( ˜ ϕ ( x )) ′ ≥ 0 . 25 .

Besides, we know that

Hence

SiLU . Note that

G Additional Lemmas

Lemma G.1. If 1 2 ≥ L 1 ≥ 1 c and L 2 ≥ L 1 -L 2 1 , we have

Proof. For function g ( x ) = x -x 2 , g ′ ( x ) = 1 -2 x . If x ≤ 1 2 , then g ( x ) is increasing. Then

Lemma G.2. Given a continuous function f ( x ) s.t. | f ( x ) - ⟨∇ f ( x ) , x ⟩| ≤ κ , then for a fixed constant r > 0 there exists C r,κ and C r s.t. for any ∥ x ∥ ≥ r ,

Now for any ∥ x ∥ > r , let y = rx ∥ x ∥ and consider g ( s ) = f ( sy ) s . Then we have

Therefore, -κ s 2 ≤ g ′ ( s ) ≤ κ s 2 . Let s = ∥ x ∥ /r ,

Therefore, f ( x ) ≤ ( C r + κ r ) · ∥ x ∥ . Similarly, we can show that -f ( x ) ≤ ( C r + κ r ) · ∥ x ∥ . Therefore, for any ∥ x ∥ ≥ r ,

Let D = max ∥ x ∥≤ r | f ( x ) | , we have for any x ,

and for any x ,

Proof. Since f is continuous, let

We have completed the proof.

Lemma G.3. Fixing c > 1 , then for every 0 < x ≤ 1 c , we have

Proof of Lemma G.3. This is equivalent to show that

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 is

This completes the proof.

Lemma G.4. Given 0 < b ≤ 1 2 and a > 0 , we have

Proof. This is equivalent to show that

This is equivalent to

Proof. Note that

Since a + b > 0 and b ≤ 1 2 , this is true.

Lemma G.5. Given c > e , we have for any x > 2 c log 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 > 2 c log c > 2 c , g ′ ( x ) < 0 . Hence, the minimal is g (2 c log c ) . Note that

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 (2 c 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 ,

Proof. Let g ( x ) = ℓ ( x + c ) -exp( -c ) ℓ ( x ) . Then, we have

Therefore, g ( x ) is monotonically decreasing. When x →∞ , we have

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 have

Then,

Therefore, c ≤ ln 3 + | x | .

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claims are accurate and reflect the contributions and scope of the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: In Sections 6 and 7, we compare our results with existing works and discuss the limitations of our approach and future directions.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All the theoretical results in the paper are accompanied by a full set of assumptions and proofs. We put all the proofs in the Appendix due to space constraints.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: All the experimental results are reproducible, and we provide all the necessary information to reproduce the results in section 5.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: Our experiments are easy to reproduce, and we provide all the necessary information in the paper. We focus on the theoretical analysis in this paper, and the code is not provided. The dataset CIFAR-10 can be found online.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All the training and test details are provided in section 5.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Our experiments are completely deterministic and it's not necessary to plot error bars. We provide the convergence plots of different models and show that the convergence rates are consistent with our theoretical results.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, 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.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [No]

Justification: We focus on the theoretical analysis in this paper, and we only provide some illustrative experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed the NeurIPS Code of Ethics and ensured that our research conforms to it.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This is 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).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

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

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We don't introduce any new assets in this paper.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/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.

14. 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.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or 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.

Peer reviews:
comments: 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.

STRENGTHS:
1. This work is the first to prove the margin improvement for large learning rates and this near-homogeneous model.
2. The authors analyze a nonlinear (near-homogeneous) model, which is closer to practice than previous works.

WEAKNESSES:
The role and effect of large learning rates, especially compared to small learning rates, may need more illustration.

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:
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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-23

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work analyzes dynamics of large step-size for GD under logistic loss for non-homogenous two layer networks. They characterize two phases, first in which empirical risk oscillates and then monotonically decreases in the second phase. Additionally, they show\n1) normalized margin grows nearly monotonically in the second phase\n2) in the oscillatory EOS phase the average empirical risk decreases under certain conditions\n3) With a larger step-size, the GD undergoing phase transition is more efficient.\n\nSTRENGTHS:\n1) The paper is well written and builds on top of exisitng works on EOS for logistic loss https://arxiv.org/pdf/2305.11788 and https://arxiv.org/pdf/2402.15926 and extends these two works from linear network to two layer non-linear network. \n2) This paper analyzes the large learning rate regime with non-homoegenous two layer networks (for logistic loss). This is a practical setting since most works study gradient flow or linear overparameterized models.\n\nWEAKNESSES:\n1) I liked the way the authors introduced the three model conditions to extend their result from Theorem-1 in https://arxiv.org/pdf/2402.15926 in terms of the three constants (Lipschitzness, smoothness and near-homogenity). However, these conditions do not hold for Relu network (assumption 1.B and 1.C breaks), which probably is most used non-linear activation out there. Infact there are a bunch of non-differentiable activation functions for which these assumptions do not hold as well as the final result. \n2) By comparing the proofs of this work with that of https://arxiv.org/pdf/2402.15926 , it seems like similar proof technique can be utilized to incorporate the three assumptions to derive the same results involving the three constants.\n\nQUESTIONS:\n1) Will the loss landscape with non-homogenous layers be the same as the linear networks case? For example, the authors show that in logistic regression, the landscape is a valley instead of a quadratic basin. The iterates converges quickly in max-margin subspace direction and the oscillations within the orthogonal direction becomes small progressively with time. Can you give an intuition of how the loss landscape may look for non-homogenous layers? \n2) Is it true that unllike MSE losses, where there is a chance of divergence for very high lrs, the logistic regression losses always converge irrespective of how large lr is used.\n3) I missed the part in the paper where the authors discuss progressive sharpening. Usually for MSE losses, PS first takes place unless the sharpness hits 2/lr after which loss oscillates around walls of basin. Does logistic regression loss do not exhibit PS? Seems like EOS starts from the offset and then oscillations stop when iterates enter a valley of flatter oscillation directions.\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies GD for nearly-1-homogeneous neural networks with large stepsizes. It provides two main results.\n \n\nThe first one describes the late, *stable phase* of training and can be seen as an extension of Lyu and Li (2020) result to the large stepsize, nearly homogeneous setting. Yet it comes with weaker conclusions, showing that if the training loss becomes smaller than some threshold at some point, then afterwards:\n     - the training loss converges to 0 at a $\\frac{1}{t}$ rate\n     - the normalized margin \"increases\", and hence converges\n\nThe second one studies the early stage, *Edge of Stability* phase for linearly separable data for a more restricted one hidden layer neural network architecture (with fixed output weights) and shows that during this early stage, the loss decreases at some rate **in average**. This then allows to provide a result illustrating the advantage of large stepsize (with an optimization error of order $\\frac{1}{T^2}$) vs small stepsize (optim error $\\frac{1}{T}$).\n\nSTRENGTHS:\nThe study of GD with large stepsizes is of high relevance. This paper puts in perspective the phenomenon of edge of stability and illustrates on a simple example how it can be beneficial to faster rates for the training loss. The claimed results seem sound and the paper is nicely written.\n\nWEAKNESSES:\nThe provided results and settings might be a bit too weak in my opinion. As a major drawback, the authors insist a lot on the importance of large step sizes for implicit bias in the paper. However, the only implicit bias result is claiming the normalized margin is increasing in the late stable phase. As commented by the authors in the paper, Lyu and Li provided a KKT point convergence result in their work, and I was hoping the same conclusion could be possible here. As a consequence, this paper mostly studies the convergence rates of the training loss, which I find less interesting from a personal point of view. \nActually, discrepancies in the normalized margins can be seen in the different experiments of the paper, but I think it should be discussed more. In the light of the literature on EoS, large stepsizes do not only lead to faster convergence rates, but also to a stronger implicit bias towards \"nice\" (e.g. sparse) features. Figure 1b) does not support this claim here, but 1c) still seems to suggest that larger stepsizes can get better test loss. In consequence, I would have also liked a similar comparison of the test accuracy in figure 2 for CIFAR-10. \n\nAdditionally, the setting focuses on nearly-1-homogeneous parameterization. While the near homogeneity assumption is mild and really nice, considering 1-homogeneous networks is very restrictive and simplistic in my opinion. Again, Lyu and Li provide a general result for $L$-homogeneous parametrizations. Having a $2$-homogeneous assumption (or even general one) would largely improve the current work. \nActually, combining this assumption with the derivative condition (Assumption 2.A) makes the considered parametrization nearly linear. As a consequence, I fear that the proof of the second main result, which depends on these assumptions, does not significantly differ from the linear regression case, from a high level, abstract point of view.\n\nQUESTIONS:\n- Is there a specific mathematical challenge (eg wrt Lyu and Li) to prove a KKT point convergence type of result?\n- Same question for extending $1$-homogeneity to $L$-homogeneity for any positive integer $L$\n- Can we say more about the normalized margin in the restricted setting of Theorem 4.1 ?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work studies the phase transition (from EoS phase to stable phase) of GD with large step sizes for training two-layer networks under logistic loss. Specifically, the authors proved the following:\n\n- If the empirical risk is below a threshold depending on the step size, GD enters a stable phase where the loss monotonically decreases and the normalized margin nearly monotonically increases.\n- For linearly separable datasets, GD with an arbitrarily large step size exits the EoS phase due to the convergence of the average loss across iterations. Moreover, a tighter bound on the phase transition time is also provided.\n- For linearly separable datasets, GD with an appropriately chosen step size achieves an accelerated convergence rate of $O(1/T^2)$.\n\nSTRENGTHS:\nThis work makes significant contributions by extending existing results. The authors proved convergence with an accelerated convergence rate in the stable phase, whereas [Wu et al. (2024)] treated the linear predictor. Additionally, the authors proved margin improvement with non-homogeneous activation functions, while previous work focused on small step sizes and homogeneous activation.\n\nWEAKNESSES:\n- Despite studying two-layer neural networks, the main theorems require linear separability of datasets, except for Theorem 2.2.\n- The paper missed some relevant works. For instance, [D. Barrett and B. Dherin] showed that gradient descent (in discrete time) optimizes the loss plus gradient norm and studied the modified ODE capturing this property. A stochastic variant of this dynamics was also studied by [Q. Li, C. Tai, and W. E]. Additionally, [M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] and [Y. Ren, C. Ma, and L. Ying] studied the benefits of large learning rates as well. Regarding optimization in the mean-field regime, [F. Chen, Z. Ren, and S. Wang] and [T. Suzuki, D. Wu, and A. Nitanda] proved the convergence of mean-field Langevin dynamics in the finite-neuron setting.\n\n[D. Barrett and B. Dherin] IMPLICIT GRADIENT REGULARIZATION. ICLR, 2021\n\n[Q. Li, C. Tai, and W. E] Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations. JMLR, 2019.\n\n[M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] SGD with Large Step Sizes Learns Sparse Features. ICML, 2023.\n\n[Y. Ren, C. Ma, and L. Ying] Understanding the Generalization Benefits of Late Learning Rate Decay. AISTATS, 2024.\n\n[F. Chen, Z. Ren, and S. Wang] Uniform-in-time propagation of chaos for mean field Langevin dynamic. 2022.\n\n[T. Suzuki, D. Wu, and A. Nitanda] Convergence of mean-field Langevin dynamics: time-space discretization, stochastic gradient, and variance reduction. NeurIPS, 2023.\n\nQUESTIONS:\n- Equation (5) in Theorem 2.2 seems to implicitly make assumptions about the data structure/distribution and the number of neurons. Can you provide any non-trivial examples other than linearly separable data? I’m wondering if this theory covers XOR or k-parity datasets, as these could be benchmarks to see the separation from the kernel regime. For instance, see the following papers:\n\n[M. Telgarsky] Feature selection and low test error in shallow low-rotation ReLU networks. ICLR, 2023.\n\n[T. Suzuki, D. Wu, K. Oko, and A. Nitanda] Feature learning via mean-field Langevin dynamics: classifying sparse parities and beyond. NeurIPS, 2023.\n- In Figure 1(c), GD with a small step size of 0.02 seems to achieve the best test accuracy at the beginning phase. What happened?\n- Can Assumption 1-C be relaxed to $\\kappa \\leq 1$?\n\nLIMITATIONS:\nThe limitations have been well addressed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}