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

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.

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

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.

Summary

This work analyzes dynamics of large step-size for GD under logistic loss for non-homogenous two layer networks. They characterize two phases, first in which empirical risk oscillates and then monotonically decreases in the second phase. Additionally, they show 1) normalized margin grows nearly monotonically in the second phase 2) in the oscillatory EOS phase the average empirical risk decreases under certain conditions 3) With a larger step-size, the GD undergoing phase transition is more efficient.

Strengths

1) The paper is well written and builds on top of exisitng works on EOS for logistic loss https://arxiv.org/pdf/2305.11788 and https://arxiv.org/pdf/2402.15926 and extends these two works from linear network to two layer non-linear network. 2) This paper analyzes the large learning rate regime with non-homoegenous two layer networks (for logistic loss). This is a practical setting since most works study gradient flow or linear overparameterized models.

Weaknesses

1) I liked the way the authors introduced the three model conditions to extend their result from Theorem-1 in https://arxiv.org/pdf/2402.15926 in terms of the three constants (Lipschitzness, smoothness and near-homogenity). However, these conditions do not hold for Relu network (assumption 1.B and 1.C breaks), which probably is most used non-linear activation out there. Infact there are a bunch of non-differentiable activation functions for which these assumptions do not hold as well as the final result. 2) By comparing the proofs of this work with that of https://arxiv.org/pdf/2402.15926 , it seems like similar proof technique can be utilized to incorporate the three assumptions to derive the same results involving the three constants.

Questions

1) Will the loss landscape with non-homogenous layers be the same as the linear networks case? For example, the authors show that in logistic regression, the landscape is a valley instead of a quadratic basin. The iterates converges quickly in max-margin subspace direction and the oscillations within the orthogonal direction becomes small progressively with time. Can you give an intuition of how the loss landscape may look for non-homogenous layers? 2) Is it true that unllike MSE losses, where there is a chance of divergence for very high lrs, the logistic regression losses always converge irrespective of how large lr is used. 3) I missed the part in the paper where the authors discuss progressive sharpening. Usually for MSE losses, PS first takes place unless the sharpness hits 2/lr after which loss oscillates around walls of basin. Does logistic regression loss do not exhibit PS? Seems like EOS starts from the offset and then oscillations stop when iterates enter a valley of flatter oscillation directions.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

As discussed the analysis breaks when considering non-differentiable activation functions like RELU. The analysis presented in the paper is an extension of previous work which I mentioned with some relaxation on the non-linearity (see the three conditions). It would be better if the authors can emphasize this limitation and point out the novelty directions in this paper.

Summary

This paper studies GD for nearly-1-homogeneous neural networks with large stepsizes. It provides two main results. The first one describes the late, *stable phase* of training and can be seen as an extension of Lyu and Li (2020) result to the large stepsize, nearly homogeneous setting. Yet it comes with weaker conclusions, showing that if the training loss becomes smaller than some threshold at some point, then afterwards: - the training loss converges to 0 at a $\frac{1}{t}$ rate - the normalized margin "increases", and hence converges The second one studies the early stage, *Edge of Stability* phase for linearly separable data for a more restricted one hidden layer neural network architecture (with fixed output weights) and shows that during this early stage, the loss decreases at some rate **in average**. This then allows to provide a result illustrating the advantage of large stepsize (with an optimization error of order $\frac{1}{T^2}$) vs small stepsize (optim error $\frac{1}{T}$).

Strengths

The study of GD with large stepsizes is of high relevance. This paper puts in perspective the phenomenon of edge of stability and illustrates on a simple example how it can be beneficial to faster rates for the training loss. The claimed results seem sound and the paper is nicely written.

Weaknesses

The provided results and settings might be a bit too weak in my opinion. As a major drawback, the authors insist a lot on the importance of large step sizes for implicit bias in the paper. However, the only implicit bias result is claiming the normalized margin is increasing in the late stable phase. As commented by the authors in the paper, Lyu and Li provided a KKT point convergence result in their work, and I was hoping the same conclusion could be possible here. As a consequence, this paper mostly studies the convergence rates of the training loss, which I find less interesting from a personal point of view. Actually, discrepancies in the normalized margins can be seen in the different experiments of the paper, but I think it should be discussed more. In the light of the literature on EoS, large stepsizes do not only lead to faster convergence rates, but also to a stronger implicit bias towards "nice" (e.g. sparse) features. Figure 1b) does not support this claim here, but 1c) still seems to suggest that larger stepsizes can get better test loss. In consequence, I would have also liked a similar comparison of the test accuracy in figure 2 for CIFAR-10. Additionally, the setting focuses on nearly-1-homogeneous parameterization. While the near homogeneity assumption is mild and really nice, considering 1-homogeneous networks is very restrictive and simplistic in my opinion. Again, Lyu and Li provide a general result for $L$-homogeneous parametrizations. Having a $2$-homogeneous assumption (or even general one) would largely improve the current work. Actually, combining this assumption with the derivative condition (Assumption 2.A) makes the considered parametrization nearly linear. As a consequence, I fear that the proof of the second main result, which depends on these assumptions, does not significantly differ from the linear regression case, from a high level, abstract point of view.

Questions

- Is there a specific mathematical challenge (eg wrt Lyu and Li) to prove a KKT point convergence type of result? - Same question for extending $1$-homogeneity to $L$-homogeneity for any positive integer $L$ - Can we say more about the normalized margin in the restricted setting of Theorem 4.1 ?

Rating

6

Confidence

4

Soundness

4

Presentation

3

Contribution

2

Limitations

N/A

Summary

This work studies the phase transition (from EoS phase to stable phase) of GD with large step sizes for training two-layer networks under logistic loss. Specifically, the authors proved the following: - If the empirical risk is below a threshold depending on the step size, GD enters a stable phase where the loss monotonically decreases and the normalized margin nearly monotonically increases. - For linearly separable datasets, GD with an arbitrarily large step size exits the EoS phase due to the convergence of the average loss across iterations. Moreover, a tighter bound on the phase transition time is also provided. - For linearly separable datasets, GD with an appropriately chosen step size achieves an accelerated convergence rate of $O(1/T^2)$.

Strengths

This work makes significant contributions by extending existing results. The authors proved convergence with an accelerated convergence rate in the stable phase, whereas [Wu et al. (2024)] treated the linear predictor. Additionally, the authors proved margin improvement with non-homogeneous activation functions, while previous work focused on small step sizes and homogeneous activation.

Weaknesses

- Despite studying two-layer neural networks, the main theorems require linear separability of datasets, except for Theorem 2.2. - The paper missed some relevant works. For instance, [D. Barrett and B. Dherin] showed that gradient descent (in discrete time) optimizes the loss plus gradient norm and studied the modified ODE capturing this property. A stochastic variant of this dynamics was also studied by [Q. Li, C. Tai, and W. E]. Additionally, [M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] and [Y. Ren, C. Ma, and L. Ying] studied the benefits of large learning rates as well. Regarding optimization in the mean-field regime, [F. Chen, Z. Ren, and S. Wang] and [T. Suzuki, D. Wu, and A. Nitanda] proved the convergence of mean-field Langevin dynamics in the finite-neuron setting. [D. Barrett and B. Dherin] IMPLICIT GRADIENT REGULARIZATION. ICLR, 2021 [Q. Li, C. Tai, and W. E] Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations. JMLR, 2019. [M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion] SGD with Large Step Sizes Learns Sparse Features. ICML, 2023. [Y. Ren, C. Ma, and L. Ying] Understanding the Generalization Benefits of Late Learning Rate Decay. AISTATS, 2024. [F. Chen, Z. Ren, and S. Wang] Uniform-in-time propagation of chaos for mean field Langevin dynamic. 2022. [T. Suzuki, D. Wu, and A. Nitanda] Convergence of mean-field Langevin dynamics: time-space discretization, stochastic gradient, and variance reduction. NeurIPS, 2023.

Questions

- Equation (5) in Theorem 2.2 seems to implicitly make assumptions about the data structure/distribution and the number of neurons. Can you provide any non-trivial examples other than linearly separable data? I’m wondering if this theory covers XOR or k-parity datasets, as these could be benchmarks to see the separation from the kernel regime. For instance, see the following papers: [M. Telgarsky] Feature selection and low test error in shallow low-rotation ReLU networks. ICLR, 2023. [T. Suzuki, D. Wu, K. Oko, and A. Nitanda] Feature learning via mean-field Langevin dynamics: classifying sparse parities and beyond. NeurIPS, 2023. - In Figure 1(c), GD with a small step size of 0.02 seems to achieve the best test accuracy at the beginning phase. What happened? - Can Assumption 1-C be relaxed to $\kappa \leq 1$?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The limitations have been well addressed.

Reviewer response-1

I thank the authors for their response. Since most of the comments seem addresses, I think Q-2 could have been addressed with more details rather than "analysis for nonlinear networks requires nontrivial efforts" and "use techniques significantly different from those of [Wu et al., 2024]". I think this question is important since the nature of results overlap with those of https://arxiv.org/pdf/2402.15926. I would suggest that the authors explain or mention how different are the proof techniques from https://arxiv.org/pdf/2402.15926 atleast briefly in the manuscript.