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

We appreciate your positive feedback. We answer your questions as follows.

---

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

A1. Corollary 4.2 and Theorem 4.3 together show a separation between large and small stepsize, where a large stepsize enables $O(1/T^2)$ loss, but small stepsizes suffer from at least a $\Omega(1/T)$ loss. This illustrates the effect of large stepsize for fast optimization.

---

Q2. What is the initial condition in Theorem 3.2?

A2. We do not have explicit assumptions on the initial condition of $w_0$. For the bound in Theorem 3.2 to be non-vacuous, $|| w_0 ||$ needs to be small compared to $\eta t$.

---

Q3. For the margin improvement part (Theorem 2.2), it would be clearer if the dependence on the learning rates in the modified margin function could be explained in the main body. Also, it would be good if there were direct comparisons between small and large learning rates results of margin improvement, e.g., different dependence on learning rates or other parameters. This may help better illustrate the advantages of using large learning rates.

A3. The modified margin is defined in Equation (10) in the appendix with explicit dependence on the stepsize. We will add this to the main paper.

The impact of the stepsize on the margin improvement is complex. Let us use $s$ to denote the starting time of the stable phase. If we assume $w_s$ is independent of the stepsize, then the loss and the parameter norm depend on the stepsize through the bounds $\Theta(1/\eta(t-s))$ and $\Theta(\log(\eta(t-s)))$, respectively. However, $w_s$ depends on the stepsize through a complex function, making it difficult to determine the quantitative effect of stepsize on margin improvement.

---

A4. Although the authors improve the theory from linear models to near-homogeneous models, it is still far from the ones used in practice. If possible, some discussions on the intuitions of more nontrivial nonlinear models may be helpful.

Q4. We believe our intuitions about the transition from an EoS phase to a stable phase extend to more general nonlinear models. Specifically, the initial EoS phase happens when GD oscillates within a sharp valley, and GD enters the stable phase when it navigates itself into a flat valley. Our theory of large stepsize is consistent with the celebrated flat minima intuition. We will add this discussion in the revision.

---

Thank you for your feedback. We address your concerns below.

---

Q1. However, these conditions do not hold for Relu network (assumption 1.B and 1.C breaks), which probably is most used non-linear activation out there. Infact there are a bunch of non-differentiable activation functions for which these assumptions do not hold as well as the final result.

A1. Note that Assumption 1.C holds for ReLU networks with $\kappa=0$ when replacing the gradients with subgradients.

Since our focus is non-homogeneous networks, we choose to work with differentiable predictors to simplify the analysis. While we do not cover ReLU activation, we have addressed many other commonly used activation functions, such as GeLU and SiLU. These two are especially interesting due to their non-homogeneity. Therefore, they cannot be covered by the prior theory. Given that our work is the first to prove margin improvement for non-homogenous predictors, we believe our contributions are very significant.

Extending our results to non-differentiable predictors using tools from [Lyu and Li, 2020] is interesting, and we will comment on this as a future direction.

---

Q2. By comparing the proofs of this work with that of https://arxiv.org/pdf/2402.15926 , it seems like similar proof technique can be utilized to incorporate the three assumptions to derive the same results involving the three constants.

A2. We emphasize that our results are for nonlinear networks, but [Wu et al., 2024] only focused on linear predictors (or networks in the NTK regime). While our EoS analysis is motivated by [Wu et al., 2024], identifying the correct set of assumptions to enable the analysis for nonlinear networks requires nontrivial efforts.

Besides, our stable phase analysis also proves margin improvement for non-homogenous predictors, which partially solves an open problem noted by [Ji and Telgarsky, 2020]. To prove this result, we use techniques significantly different from those of [Wu et al., 2024].

---

Q3. Will the loss landscape with non-homogenous layers be the same as the linear networks case? For example, the authors show that in logistic regression, the landscape is a valley instead of a quadratic basin. The iterates converges quickly in max-margin subspace direction and the oscillations within the orthogonal direction becomes small progressively with time. Can you give an intuition of how the loss landscape may look for non-homogenous layers?

A3. No, the loss landscape for non-homogeneous networks is significantly different from that of logistic regression. Note that the formal is non-convex while the latter is convex. It is difficult to develop geometric intuitions about the non-convex landscape of non-homogeneous networks. This difficulty also demonstrates the significance of our contributions to establishing margin improvement and EoS theory for non-homogeneous networks.

---

Q4. Is it true that unllike MSE losses, where there is a chance of divergence for very high lrs, the logistic regression losses always converge irrespective of how large lr is used.

A4. This is a good question and we do not know the answer. In practice, the loss can diverge when training deep networks with very large stepsizes, even under logistic loss. However, we are not sure if this is a numerical issue or if there is a provable upper bound for the convergent stepsize.

---

Q5. I missed the part in the paper where the authors discuss progressive sharpening. Usually for MSE losses, PS first takes place unless the sharpness hits 2/lr after which loss oscillates around walls of basin. Does logistic regression loss do not exhibit PS? Seems like EOS starts from the offset and then oscillations stop when iterates enter a valley of flatter oscillation directions.

A5. Empirically, progressive sharpening also happens under logistic loss, as shown by [Cohen et al., 2020]. However, this is beyond the scope of this work, which focuses on the implicit bias and convergence of large stepsize GD.

---

Thank you for your feedback. We address your comments below.

---

Q1. Is there a specific mathematical challenge (eg wrt Lyu and Li) to prove a KKT point convergence type of result?

A1. Good question. The key difficulty is that KKT points with respect to an optimization problem are not even well defined for non-homogenous predictors, so extending the KKT point convergence analysis Lyu and Li is challenging.

Since our nearly homogeneous function $f(w;x)$ is asymptotically homogeneous, we conjecture that, under suitable conditions, GD may converge to KKT points given by a proxy homogeneous function

$$\tilde f(w;x) := \left [ \lim_{t \to \infty} \frac{f(t w;x)}{t} \right] \cdot ||w||.$$

There are still two obstacles to proving the above conjecture.

1. Verifying the (sub)differentiability of $\tilde f(w;x)$ is nontrivial.
2. The GD dynamic for learning $f$ is different to that for learning $\tilde f$. Specifically, there are examples such that $|f(w)-\tilde f(w)|$ is uniformly lower bounded by a fixed constant. Then the GD dynamic for $f$ is away from that for $\tilde f$, and might not satisfy the corresponding $(\epsilon, \delta)$ KKT conditions in [Lyu and Li].

We believe additional regularity assumptions about $f$ are needed to show KKT convergence. Nonetheless, we think this is a great question, and we will discuss it in the revision.

---

Q2. Is there a specific mathematical challenge (eg wrt Lyu and Li) to extending 1-homogeneity to L-homogeneity for any positive integer?

A2. As mentioned in Line 300, this is an avenue left for future exploration. Extending our results from near 1-homogeneity to near L-homogeneity will require non-trivial modifications to both the margin functions and the stable phase conditions outlined in Theorem 2.2. Furthermore, verifying that the proposed stable phase conditions can be satisfied (even in special cases) could be challenging. We will discuss this in the revision. We think our margin improvement results in this work, which extends prior results from 1-homogeneity to near 1-homogeneity, is already a significant step toward addressing the general case.

---

Q3. Can we say more about the normalized margin in the restricted setting of Theorem 4.1?

A3. Good question. Although we consider linearly separable data, the predictor is nonlinear and, therefore, can achieve a normalized margin larger than the maximum $\ell_2$-margin. Additionally, we observe that in practice, while the normalized margin tends to be positive and increasing, the normalized margins for individual neurons can stay negative (please check Figure 2 in the pdf). A full characterization of the normalized margin is technically challenging even for linearly separable data. We will comment on this as a future direction in the revision.

---

Q4. The provided results and settings might be a bit too weak in my opinion… the only implicit bias result is claiming the normalized margin is increasing in the late stable phase…. In the light of the literature on EoS, large stepsizes do not only lead to faster convergence rates, but also to a stronger implicit bias towards "nice" (e.g. sparse) features. Figure 1b) does not support this claim here, but 1c) still seems to suggest that larger stepsizes can get better test loss. I would have also liked a similar comparison of the test accuracy in figure 2 for CIFAR-10.

A4. We respectfully disagree that our results and settings are weak. In fact, extending the implicit bias results for homogeneous predictors to general non-homogeneous predictors is an open problem listed in [Ji and Telgarsky, 2020]. Our work, by establishing the margin improvement results to general near 1-homogeneous predictors, partially solves this open problem. Furthermore, we also extend prior neural network analysis for GD with small or infinitesimal stepsizes to an arbitrary large stepsize, which is more relevant to practice. We believe the contributions of this work are already significant.

While understanding the generalization benefits of large stepsizes is very interesting, this question is beyond the scope of the current paper, which focuses on margin improvement and fast optimization.

Please see the attached pdf for Figure 1 reporting the test accuracy for CIFAR-10. We will include it in the revision.

---

Q5. ….Actually, combining this assumption with the derivative condition (Assumption 2.A) makes the considered parametrization nearly linear….

A5. We would like to point out that Assumption 2.A is a sufficient technical condition to enable the EoS phase analysis. However, we only need Assumption 1.A for Theorem 2.2, which does not assume a lower bound on the derivative and allows the predictor to be highly non-linear.

---

Thank you for your support.

---

Q1. Despite studying two-layer neural networks, the main theorems require linear separability of datasets, except for Theorem 2.2.

A1. We acknowledge that linear separability is a strong assumption. We use this mainly as a sufficient condition for two-layer neural networks, regardless of width, to reach the initial bound of the stable phase when employing large stepsizes. Additionally, our stable phase results do not need this assumption. We believe the linear separability condition can be relaxed, which is left for future work. We will comment on this more in the revision.

---

Q2. The paper missed some relevant works. For instance, [D. Barrett and B. Dherin] showed that gradient descent (in discrete time) optimizes the loss plus gradient norm and studied the modified ODE capturing this property. A stochastic variant of this dynamics was also studied by [Q. Li, C. Tai, and W. E]. Additionally, [M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] and [Y. Ren, C. Ma, and L. Ying] studied the benefits of large learning rates as well. Regarding optimization in the mean-field regime, [F. Chen, Z. Ren, and S. Wang] and [T. Suzuki, D. Wu, and A. Nitanda] proved the convergence of mean-field Langevin dynamics in the finite-neuron setting.

A2. Thank you for bringing our attention to these works. We will cite and discuss them in detail. Specifically, the stepsize considered by [D. Barrett and B. Dherin] and [Q. Li, C. Tai, and W. E] are small or even infinitesimal, whereas our stepsize is large and causes EoS. The work by [Y. Ren, C. Ma, and L. Ying] studied special linear networks under MSE loss, while we focus on two-layer networks under logistic loss. The work by [M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] studied the effect of large stepsize through experiments while we take theoretical approaches. The works by [F. Chen, Z. Ren, and S. Wang] and [T. Suzuki, D. Wu, and A. Nitanda] demonstrated the convergence of mean-field Langevin dynamics for two-layer networks with a finite but large number of neurons. In comparison, although we adopt the parameter scaling from mean-field theory, we use a different proof technique that allows any number of neurons.

---

Q3. Equation (5) in Theorem 2.2 seems to implicitly make assumptions about the data structure/distribution and the number of neurons. Can you provide any non-trivial examples other than linearly separable data? I’m wondering if this theory covers XOR or k-parity datasets, as these could be benchmarks to see the separation from the kernel regime.

A3. This is a good question. In fact, we only need to ensure that the dataset is realizable, that is, there exists a parameter direction $W$ such that $\inf y_i f(a W; x_i) \to \infty$ as $a\to \infty$. This indicates that the loss can be made arbitrarily small. Since XOR or k-parity datasets can be realized by two-layer networks (with Softplus, for instance), Equation (5) can be satisfied. We will discuss this in the revision.

---

Q4. In Figure 1(c), GD with a small stepsize of 0.02 seems to achieve the best test accuracy at the beginning phase. What happened?

A4. You are correct that GD with a small stepsize achieves the best accuracy early (after about 20000 iterations). We think this is because of a mild overfitting since we only use a small dataset of 1,979 samples, while the MLP has about 30,000 parameters.

---

Q5. Can assumption 1-C be relaxed to $\kappa\le 1$?

A5. In Assumption 1-C, $\kappa$ can be any fixed nonnegative number.

---