Implicit Bias of Gradient Descent for Two-layer ReLU and Leaky ReLU Networks on Nearly-orthogonal Data

Summary

The paper studies convergence and implicit bias in two-layer (leaky) ReLU networks trained with gradient descent. They focus on the case where the initialization scale and the step size of GD are sufficiently small, and the training data are nearly orthogonal. They show two main results: - For leaky ReLU networks, they prove convergence to stable-rank 1 of the weight matrices (they consider the matrix of the positive neurons and the matrix of the negative neurons separately), and also give a convergence rate analysis; - For ReLU networks they bound the stable ranks by a constant, which in some special cases can be roughly 2. They analyze the convergence rate also here.

Strengths

I think that it is a nice paper. The near-orthgonality assumption is quite strong, but I believe that it is a natural assumption, and that understanding the behavior of GD in this setting is helpful. For this reason, this assumption appears (either explicitly or implicitly) in several works in recent years. This work extends the results from Frei et al. 2022b in several ways: - Frei et al. 2022b showed convergence to rank 1 of the matrices when using gradient flow and constant stable rank when using GD, and here it is extended to stable rank 1 using GD; - Frei et al. 2022b showed convergence to a constant stable rank when using GD with smooth leaky ReLU activation, and here it is extendedto leaky ReLU and to ReLU; - The convergence rate here is tighter than Frei et al. 2022b. I think that the question considered in the paper is interesting and the contribution is significant and non-trivial. Also, the paper is well-written.

Weaknesses

I don’t see major weaknesses. Below I discuss some comments/questions that I have. - In the introduction the authors argue that “This finding suggests that ReLU networks possess superior learning ability compared to leaky ReLU networks”. I don’t see how this claim follows from the results. This sentence appears after mentioning the example where the stable rank becomes approximately 2 in ReLU networks, but this occurs also in leaky ReLU networks for this example. - Related to the above comment: Relating low stable rank to learning ability is not a trivial argument. While the exact rank can be used for bounding the complexity of the hypothesis class, this is not true in general for the stable rank. So I don’t see a direct argument that gives learning guarantees by only using the stable rank. (Still, I think that the stable rank is a natural notion and hence studying it is of interest). - The experiments on MNIST are conceptually similar to the experiments on CIFAR from Frei et al. 2022b. Please explain what is the contribution of this experiment compared to this prior work. - Typos: - In Lemma 5.4 (line 215): should be $i,i’ \in [n]$ instead of $r,r’ \in [m]$. - Line 286: “RelU” - Line 759 (in Lemma F.2): - In the first property should be $\rho_{j,r,i}$ - The second property should be without the \bar{} - All conditions should consiter $w^{(0)}$ rather than $w^{(t)}$ - Last equation in page 42: should be $0_{m \times(d-n)}$. Also in the equation before that.

Questions

See the “weaknesses” section.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The authors discussed the limitations.

Summary

The paper investigates implicit bias, especially in terms of the stable rank of the weight matrix (which is the square of the ratio between the Frobenius and spectral norms), of gradient descent on the logistic loss for binary classification tasks and two-layer neural networks, where the first layer is trained and the second layer is fixed, the activation function is either ReLU or Leaky ReLU, and the training data are nearly orthogonal. The main results are that, provided the learning rate and the initialisation scale are sufficiently small, as the number of training steps tends to infinity, the stable rank tends to 1 in the Leaky ReLU case, and is bounded by a constant in the ReLU case. The authors also prove upper bounds on the convergence rate of the empirical loss, and tight bounds on the convergence rate of the norms of individual neurons in the Leaky ReLU case. The theoretical results are supplemented by experiments on synthetic data generated by adding Gaussian noise, and on MNIST.

Strengths

The paper tackles one of the greatest current challenges in the theory of learning, namely understanding the implicit bias of gradient-based algorithms and its consequences for generalisation. In contrast to most related works, it considers gradient descent rather than gradient flow, and is thus more readily applicable to practical settings. The activation functions that it focuses on, namely ReLU and Leaky ReLU, are among the most popular. The theoretical results are proved in detail in the extensive appendix, and sufficient details as well as the code for the experiments are provided.

Weaknesses

The paper is a follow-up to "Implicit Bias in Leaky ReLU Networks Trained on High-Dimensional Data" by Frei, Vardi, Bartlett, Srebro, and Hu (ICLR 2023), and shares most of the assumptions with the gradient descent part of that work, with the important difference that in the present paper the activation functions are ReLU and Leaky ReLU rather than a smooth variant of Leaky ReLU as in that work. It is not very clear to me what difficulties exactly needed to be overcome to step up from that previous work to this paper, or in other words to what extent and in what ways the present paper is not a straightforward increment. Another aspect that is similarly unclear to me is how essentially different are the ReLU and Leaky ReLU cases. The theorems and proofs for them are presented in the main and in the appendix mostly separately, yet seem to follow similar patterns, so it would be good to clarify where the main underlying differences are. A minor comment is that there is a typo saying that the signal-noise decomposition technique is introduced in the Preliminaries.

Questions

What can you say about what happens if the second layer is trained as well, or if the network is deeper?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

Training only the first layer significantly simplifies the dynamics.

Summary

This paper studies training two-layer leaky-ReLU/ReLU networks with gradient descent. The authors show that for approximately orthogonal data, GD converges and the weight matrices in the leaky-ReLU network have a stable rank of one asymptotically under GD, while those in the ReLU network can have a stable rank larger than one.

Strengths

1. The convergence analysis of GD on two-layer neural networks seems novel 2. The characterization of the stable rank of the weight matrices is interesting

Weaknesses

1. There is little discussion on the upper bound $c$ in Theorem 4.3. How large it can be? How does it depend on the training data? 2. For leaky-ReLU, the theorem only says it has a stable-rank one, meaning $W_j$ is approximately rank-one, is there a characterization of the left and right dominant singular vectors?

Questions

See "Weakness"

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

yes

Summary

This study analyzes two-layer ReLU and leaky ReLU networks trained by gradient descent in a classification problem. It is shown that when training data is mutually nearly-orthogonal, by taking an initialization scale and a learning rate sufficiently small, a leaky ReLU network converges to a parameter with stable rank 1. They also show that a ReLU network converges to a parameter with a bounded stable rank, and it will equal 2 when the input is completely orthogonal. A similar analysis has already been provided by Frei et al. 2022 (with gradient flow and Leaky ReLU), this work takes another approach apart from the KKT condition, and a possibly faster convergence rate is obtained.

Strengths

- Frei et al. (2022) have already shown a certain result on this topic, but it requires several non-practical modifications of the algorithm (such as considering gradient flow or smoothed leaky ReLU). This study overcomes this shortcoming with a novel approach and obtains further results (extension for the ReLU activation and gradient descent without smoothing), which has an important novelty.

Weaknesses

- The near-orthogonality assumption $R_{\min}^2\ge CR^2\gamma^{-4}np$ seems to be strong (while Frei et al. (2022) also impose the same condition, but I realize that they assume high-dimensional settings). - It will be more interesting if the authors can characterize the direction to which the weight matrix converges. - The authors refer to the works on benign overfitting in the related work section, but the connection of this study to them needs to be made clearer. Does this study focus on high-dimensional settings? Or is the double descent curve provably obtained by this study? - Some terms are used without definition, such as $W_{y_i}^{(t)}$ in (5.8) and $w_{y_i,r}^{(0}$ in Lemma 5.5, and so on.

Questions

- The authors conduct numerical experiments and observe the stable rank of trained networks, but I'm wondering about it precisely reflects the theoretical result. That is, whether the stable rank of the leaky ReLU network with large $\gamma$ converges to 1 or not, and the stable rank of the ReLU network could converge to 2 by using fully-orthogonal training data.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A