Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: summary: This paper studies the training dynamics of GD and SGD in mildly over-parameterized neural networks. Specifically, the major result focuses on the early-stage training dynamics and proves that the loss decreases by a significant amount in the early iterations. Global convergence result is also provided.
strengths_and_weaknesses: The paper is technically solid. I think the most interesting result is the early-stage result, as it presents a concrete hitting time when the loss descends to a relatively low value. In particular, the hitting time is proved to be small, which matches the empirical observation that the training loss descends rapidly in the first few epochs. Therefore, this early-stage result might be able to explain practical training. The partition technique also looks interesting and may help in understanding the training dynamics of GD in neural nets.

There are still some weak points in this paper.

I skim the appendix and believe the proof is correct. However, the author seems to assume specific values of $\eta$, $\kappa$ and $m$ in the proof. It is better to keep $\eta$, $\kappa$ and $m$ in the complete proof to better observe the relationship between the loss bound and these parameters.

The assumption is still a bit strong. Specifically, the initialization is assumed to be close to 0. This assumption may be not desirable in the following two aspects: 1. It may not characterize the training dynamics in the global region; 2. It may not match practice.

The settings are not unified in binary and multi-class classification. Specifically, for binary classification, the paper uses the ReLU network without bias; whereas for multi-class classification, the paper uses the ReLU network with bias. It would be better if the settings are changed to the same.

The comparison with the results in late-stage training (Line 50, 185) looks a bit strange. In my humble point of view, these results focus on the generalization performance of GD in neural networks and are orthogonal to the topic of this paper.

The writing of this paper could be improved. The partition names in Lemma 6.2 look confusing as “false-dead” neurons are truly dead in the sense that they remain dead, and “true-dead” neurons are not dead as they can turn into true-living. There are also some grammar mistakes in the paper, especially in the singular/plural form and the usage of articles (Line 41, 42, 55, 74, etc.).

questions: Line 57: Is it "mildly overparameterized" instead of "mildly parameterized"?
limitations: Yes.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: This paper investigates the dynamics of ReLU neural networks in the initial portion of training. By using the discontinuity of the ReLU activation function on the hidden layer as well as a structured dataset a partitioning of the hidden layer neurons is defined for each datapoint. Specifically hidden neurons are grouped into one of four partitions based on whether their input-to-hidden (hidden-to-output) weights have positive/negative dot product with the input (output). From this structured initial setup it is possible to partition the initial hidden neurons based on the initial weight distributions. This then allows for bounds on the network error to be defined since the dynamics of the partitioned subnetworks are more deterministic and have discernible properties (such as most neurons remaining in their original partition). Thus, the primary results of this work are theoretical statements on the lower bound which the error will decrease by in the early stages of learning and a global convergence theorem for the structured dataset.  
strengths_and_weaknesses: # Strengths
## Originality
To my knowledge this is the first work which has theoretically used neuron partitioning, however, as the paper acknowledges, this has previously been done empirically. Focusing on the early stages of learning as opposed to the more prominent approach of investigating the late stages of learning also adds to the originality. 

## Quality
The paper is of a high quality and well structures. The research itself is focused in its goal and the methodology appears appropriate. The only figure in the paper is clear and helpful.

## Clarity
The paper is clear and well written. Notation is intuitive, used sparingly and used consistently which aids in the readability. Figure 1 adds to the clarity of the work and displays a helpful setting that demonstrates how Lemma 6.2 can be consistent for different datapoints.

## Significance
This appears to be the most difficult part to assess and I do have some reservations (see below). However, a tighter bound on the gradient norm is given and the theoretical technique of partitioning the hidden neurons could potentially yield future insight. For me the rigor of this work is what adds to the significance substantially. Having gone through Appendix A quite carefully and Appendix C (it seems the proof technique is similar for all Appendices) the proofs seem correct and rigorous.

# Weaknesses
## Originaity
This work does bear some resemblance to prior theoretical work on gated neural networks, see [1] for example. I think this should be acknowledged.

## Quality
I have 3 primary concern from the main text of this work
1. This work could benefit greatly from a simulation of the derived equations. While I think purely theoretical work is certainly acceptable, the theory should still be checked against even a toy example. The XOR dataset [2] for example seems like it would be appropriate and provide a sense of how tight the derived bounds really are. Adding something like this is my main recommendation and its absence my main reservation.
2. On line 98 and 99 the "less than or equal to some multiple" operator $\lesssim$ is defined. This operator appears very loose since $a \lesssim b$ translates to $a \leq cb$ for some positive c value. This appears to be the same as saying $a \leq \infty$ assuming b is positive. As far as I can see it is used sparingly. However, it is used early on in the main text to define the learning rate which is obviously an important variable and impacts other variables such as $\kappa$. It is also used for Lemma 6.4 which is used for the global convergence Theorem. What is the effect of this on the tightness of the bounds and perhaps it is necessary to address this.
3. The learning rate is another point of concern in that some learning rates presented should be unstable, yet no mention is made of this in the text. Specifically if the learning rate must be less than a particular ratio of the top eigenvalue of the hessian to avoid divergence [3]. Such a limitation is not place when defining the learning rate (as mentioned for the second point above) or for Theorem 5.2 where the learning rate is inversely proportional to the loss (and will tend to infinty as a result). I don't think ignoring stability in this case is an issue, however, I again think it should be mentioned.

## Significance
The three points under quality above also limit the  significance, or rather limit my ability to determine this work's significance. The main issue again being that it is unclear how tight the bounds really are or how applicable they are without some form of simulation being run. If not such simulation fits all of the assumptions then this would need to be acknowledged. If this is the case I do not view it as grounds for rejection but it is important for a reader to know. That said, this work seems to be a good addition to the theoretical literature and I think the strength outweigh the weaknesses for its significance.

[1] Lakshminarayanan, Chandrashekar, and Amit Vikram Singh. "Deep Gated Networks: A framework to understand training and generalisation in deep learning." arXiv preprint arXiv:2002.03996 (2020).

[2] Brutzkus, Alon, and Amir Globerson. "Why do larger models generalize better? A theoretical perspective via the XOR problem." International Conference on Machine Learning. PMLR, 2019.

[3] Cohen, Jeremy M., et al. "Gradient descent on neural networks typically occurs at the edge of stability." arXiv preprint arXiv:2103.00065 (2021).

questions: ## Main Text
I have the following questions or minor points which could be clarified (in addition to the questions which fit in with the Strengths/Weaknesses above):
* Line 55: two full-stops after "neuron"
* Line 58: It is mentioned that the loss landscape will be complicated at initialization, however, with small initial weights this does not appear to be true as the model will be sitting very close to a smooth saddle point at 0.
* Line 95: It says capital letters are used to denote vectors but the example shows a lower case bold letter for a vector (which is used in the rest of the work)
* Line 131 to 134: As a theoretical tool using the labels to preprocess data is fine (and I acknowledge that this assumption is not new to this work), but it limits the scope of the theory as it then cannot apply to unlabelled data or test data. Again I think this could be mentioned.
* Assumption 4.5: it never becomes apparent to me why this assumption is necessary or different from Assumption 4.1. I would just like to understand that better
* Line 169: Can remove "the"
* Line 216: I assume the $\lambda$'s are the singular values but we should be told what they are.
* For Theorem 5.3: Are only the input-to-hidden weights being trained?

## Appendix (optional)
* Line 519: $0 \leq s \leq t+1$ is stated twice
* Line 520: Where does the $\gamma_1$ variable come from or how is it determined?
* Line 539: what is $P^\infty$, if it is the infinite width limit for the hidden layer then does this fit in with the notion of this theory being for mildly over-parametrized models?
* Line 541: How does the following line apply from the union bound? Is it the union bound over the hidden neuron distributions?
* Line 547: Is this mean to say "union bound"?
* Line 560: Is the final piece of this line guaranteed since from the line before the $-1-f(x,\theta)$ make it seem like the entire second part of the expression could be negative?

As a guide, if points 2 and 3 from the points under weaknesses for Quality are addressed I would be inclined to increase my score by 1. If point 1 is addressed and the results are compelling I would increase my score by 1.
limitations: There appear to be a couple limitations which could be acknowledged that I have pointed out above. If I am not mistaken then I think they should be mentioned.
review_title: null

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: summary: This paper studies the early stage of convergence for moderate width networks optimized via the squared loss and logistic-style loss functions.  More precisely, they demonstrate that the loss will descend by an $\Omega(1)$ quantity in the early stage of training.  If the data are very well-separated, for logistic-style loss functions they can demonstrate global convergence.  A contribution is that they only require logarithmic width for the network.  As a matter of technique the novelty is that rather than prove NTK deviation bounds they control the NTK gram matrix via fine grained analysis of the neurons.  While previous works such as [1] and [2] use overparameterization to demonstrate that the activation patterns do not change much, the work under consideration does not rely on overparameterization and performs a more refined analysis of neurons based on true/false-alive/dead pairings.

References:

[1] Gradient Descent Provably Optimizes Over-parameterized Neural Networks by Du et al.

[2] Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks by Arora et al.
strengths_and_weaknesses: 
The strength of the paper is that it only requires logarithmically wide networks and it provides a fine-grained analysis of the neurons that is novel as far as I can tell.  Furthermore the writing and presentation is easy to follow.  However, I have some concerns about this work.  I would like the authors to provide a more detailed explanation about how their analysis is able to prove results that cannot be proven using existing tools from NTK theory. If one is only interested in analyzing the early stage of training one does not need overparameterization to perform standard NTK analysis.  Up to a $O(1)$ stopping time, parameter deviations are $O(1)$ and thus NTK deviations are $O(1/\sqrt{m})$.  Thus if one is only interested in the early-time dynamics the width of the network can be taken independent of $n$.  As far as the statements for global convergence, under Assumption 4.1 the data are linearly separable.  When the data are linearly separable with positive margin, the target function is separable with positive margin using the NTK model see [3].  For this case [3] demonstrated that polylogarithmic width suffices for the logistic loss to achieve global convergence.  To be fair in [3] they only train the hidden layer, however due to $\kappa$ being small in the present work the gradients of the hidden layer will typically be larger than the gradients for the outer layer.  In this case the changes in the hidden layer should be more significant than the changes in the outer layer.  Thus I do not see this being too essential here.  Also in Theorem 5.3 the authors only train the hidden layer which allows a more direct comparison.

References:

[3] Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks by Ji and Telgarsky https://arxiv.org/pdf/1909.12292.pdf
questions: Most of my questions were outlined in the “Strengths and Weaknesses” section, however I will reiterate these here as well as introduce some new questions.

Question 1: Could the authors provide a more detailed discussion of how their analysis of the early stage of training is able to improve over what more standard NTK analysis would give.

Question 2: 2nd inequality of line 570 could be explained better.

Question 3: In Lemma A.13 the authors upper bound the gradient deviations using the loss Hessian.  However, for ReLU the loss Hessian does not exist in the proper sense.  I believe this can be overcome using a different argument however this discrepancy must be addressed or at least explained.

Question 4: Could the authors provide a more detailed discussion of the scale of the initialization.  If $\kappa$ is small then the gradients of the inner layer weights will dominate over the gradients of the outer layer weights (due to properties of ReLU).  This is not bad necessarily however this should be explained to the reader.  I believe this is essential to the proofs.  For example the authors study the stopping time $T$ up to which the signs of the outer layer weights haven’t changed.  The size of $\kappa$ will be quite relevant for determining how quickly the outer layer weights change in sign.

Question 5: Could the authors provide a more detailed comparison of their work and [3] which addresses some of my concerns outlined in the “Strengths and Weaknesses” section.

References:


[3] Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks by Ji and Telgarsky https://arxiv.org/pdf/1909.12292.pdf
limitations: The main contribution of this work as I see it is to provide a fine grained analysis of the neurons for the shallow ReLU network.  This is indeed interesting, although it will be difficult to generalize to multiple layers or to activation functions other than ReLU.

Based on my concerns in the "Strengths and Weaknesses" section, I do not view the present version of this manuscript to be ready for publication in this conference.  If the authors provide strong responses to my concerns to show that I am mistaken or missed crucial points I will consider raising my recommendation.

**Edit after rebuttal: I upgraded my recommendation from a 3 to a 5 based on the edits and responses from the author.  Please see my comments below for my reasoning.**
review_title: null

final_score: 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.
review_content: summary: This paper studies the convergence of GD and SGD when training mildly parameterized neural networks starting from random initialization. The authors prove an "early stage convergence" that the loss is decreased by a significant amount in the early stage of the training quickly. 
strengths_and_weaknesses: The paper is well-written, technique sound, and studies important problems of neural network training process, one is the fast convergence at the early training stage, the other is the global convergence of mildly-overparameterized neural networks. However, some assumptions may be not so realistic, which makes reader hard to connect the results and practical training process.
questions: Assumption 4.1 assumes that close data points have same output, this also implies the training data has significant low-frequency components. According the frequency principle (https://arxiv.org/abs/1901.06523), low frequency would converge at the beginning and fast, Theorem 4.2 can be intuitively understood. For Assumption 4.5, it constraints the input but not the output. That means Theorem 4.6 applies for any training data. However, in practical training, if the training data is dominated by very high frequency components, the loss may not decay or decay extremely slow. Does this example show that the \Omega(1) bound in theorem 4.6 may be loose?

The model in Theorem 4.2 lacks of universal approximation ability (without bias), how does this result help understand a model with bias. 

A loss that satisfies Assumption 4.4 may not a common case in deep learning, such as MSE and cross-entropy. Can authors explain more about how Theorem 4.6 depends on the Assumption 4.4? 

I am still confusing about how Theorem 4.2 and 4.6 help understanding the training process. For example, can these two theorems tell us in which situation there is an early-stage convergence and there is not, or to achieve a fast convergence at the initial stage, how should we design the network.

Regarding the global convergence part, the results heavily relay on the Assumption 4.1 and model (4), which seems not very realistic. Can authors discuss more on how these results can help understand the practical neural network?

limitations: The theoretical results in this paper are nice and sound solid, however, I suggest the author should make the limitation clear in abstract and introduction, which are now quite misleading. The informal statements of theorem ignore very important assumptions, which will mislead readers by regarding the results as very general results. The title and the abstract should explicitly point out the paper studies two-layer neural network at least. These limitations at least include the model, loss and data.
review_title: null

Meta-review: metareview: There was an extensive discussion of the article concluding in reviewer recommendations weak accept, borderline accept, and accept. Although there were some reservations, particularly about the assumptions on the data, model and loss, the reviewers found that the article is well-written, technically sound, that it studies important problems of neural network training process, and is valuable for providing a novel analysis. Based on these merits, I am recommending accept. However, I will ask that the authors carefully consider the extensive feedback in the preparation of the final manuscript, particularly the comments concerning the presentation and discussion of the assumptions and limitations, and also carefully work on the improvements discussed during the discussion period, particularly the technical detail of the non Lipschitz gradient, as well as the promised additions, such as the proofs for the quadratic loss in the multi-class setting. 

confidence: Certain
Decision: Accept