Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This submission studies the gradient descent/flow dynamics of a two-layer neural network with ReLU activation function and losses with exponential tail. The authors presented the following results: (i) an improved perceptron analysis based on the NTK margin that allows the parameters to move away from the NTK regime, and it is shown that a polynomial-width neural network achieves low test error; when specialized to the 2-parity problem, this gives learning guarantees that matches kernel methods. 
(ii) under a neural collapse assumption, an exponential-width neural network with exponential loss can achieve large margin and low test error. (iii) in a scalar gradient flow where the neurons cannot rotate, it is shown that an exponential-width neural network can find the global max-margin solution (up to constant); for the 2-parity problem, this implies superiority over kernel methods. 
strength_and_weaknesses: I find this submission fun to read. Despite the highly technical/theoretical content, the authors provide nice intuitive explanations of the analysis in the main text (e.g. on the construction of potential function). Also, the limitations of each approach and comparison to prior results are carefully discussed. I believe the following contributions are meaningful: 

- The NTK margin analysis allows the parameters to move up to $\mathcal{O}(\sqrt{m})$ and improves the required width in terms of the NTK margin. It is also interesting that the margin-based Rademacher complexity bound has no dependence on the network width (due to manipulation of the homogeneity property).  
- To my knowledge the scalar gradient flow analysis is novel. It outlines a new assumption (low rotation) that leads to global margin maximization, which differs from prior results that assume noisy dynamics or dual convergence. Moreover, this low-rotation assumption seems to hold empirically when the network width is large. 

The weaknesses of this submission are also clear (as acknowledged by the authors). If we focus on the sample complexity of learning the 2-sparse parity function, then the improvement over existing results is not very significant. While the perceptron analysis is not an NTK proof, the resulting sample complexity does not show any advantage over the NTK model. As for the scalar gradient flow, the required model width still exponential even under this rotation-free parameterization. 

I have the following questions/comments. 

1. The authors used the 2-sparse parity problem as an illustrative example throughout the paper. What is the motivation of using this particular problem? How would the margin values change if we generalize to $k$-parity? And is the sample complexity still better than kernel methods? 

2. Is there an intuitive explanation of why the large-width analysis can only handle the exponential loss? 

3. Is the empirically-observed low-rotation property specific to the ReLU (or homogeneous) activation function? 

4. A few more related works: the learning of parity-like functions using two-layer neural networks has been studied in [1] [2]. The two-phase algorithm (learning the features and then the second layer) has been analyzed in the regression setting in [3] [4].  
[1] Refinetti et al. Classifying high-dimensional gaussian mixtures: where kernel methods fail and neural networks succeed.  
[2] Frei et al. Random feature amplification: Feature learning and generalization in neural networks.  
[3] Ba et al. High-dimensional asymptotics of feature learning: how one gradient step improves the representation.  
[4] Bietti et al. Learning single-index models with shallow neural networks.  
clarity,_quality,_novelty_and_reproducibility: See strength and weaknesses. 
summary_of_the_review: I am not very familiar with the max-margin literature, and I did not read the Appendix in details. 
However, I believe the results presented in this submission are relevant to the ICLR community.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper considers shallow neural networks trained with gradient methods (stochastic gradient descent or gradient flow) under logistic or exponential loss. The authors prove that a low test error can be achieved in a regime that still exhibits low rotations in the weights, but is capable of exiting the NTK regime (in which the weights do not move much from their initialization).

More specifically, the authors consider networks with moderate width (order of $d^2$) and prove that a margin scaling linearly with $\gamma_{\rm ntk}$ can be achieved by either stochastic gradient descent (Theorem 2.3) or gradient flow (Theorem 2.1). This improves upon the width requirement in an earlier work by Ji & Telgarsky, (2020b) and upon the number of samples needed by an earlier work by Barak et al., (2022) -- although reaching the Pareto frontier of width, samples and computation remains a (possibly rather difficult) open problem.

Next, the authors consider a setting close to the 'neural collapse' (NC) popularized by the recent paper by Papyan et al., (2020). Here, the networks widths are impractical (exponential in $d$), but the authors can show a margin scaling linearly with the *global* margin $\gamma_{\rm gl}$ -- as opposed to the previous results scaling with $\gamma_{\rm ntk}$. These results also require either an assumption on a NC property of the dataset (Theorem 3.2) or on the inability of the inner weights to rotate (Theorem 3.3).
strength_and_weaknesses: Strengths:

(1) Theorem 3.3 improves in the network width over (Ji & Telgarsky, 2020b) and in the number of samples over (Barak et al., 2022). 

(2) Interesting technical innovations over existing work.  

(3) The assumption on low rotation, although strong, does allow for stronger results (achieving the global margin) and seems to be satisfied experimentally.

(4) [Minor, but still nice to see...] Well written paper with the various assumptions properly discussed, the related work discussed in details, and interesting open problems. 

Weaknesses: 

(1) Model still somewhat artificial.

(2) The optimality of the results remains unclear. Theorem 2.3 leads to a bound on the network width of order $d^2$. This follows from having the width scaling as $1/\gamma_{\rm ntk}^2$ and then showing that $\gamma_{\rm ntk}$ scales at least linearly in $1/d$. Is any of these two bounds optimal (in some sense)? For the regression setting, the minimum width has to scale linearly with the number of parameters (this follows from counting degrees of freedom, or from VC-dimension bounds). However, for classification, it is not entirely clear what the minimum width should be. I would encourage the authors to discuss in more detail this point, which is only briefly mentioned as the open problem of understanding 'the Pareto frontier of width, samples, and computation'. 

(3) In the statement of Theorem 2.1, the authors claim that there exist $t$ with $||W_t-W_0||=\gamma_{\rm ntk}\sqrt{m}/32$. Did the authors intend to put an upper bound here? 

(4) Is the probability bound of Theorem 2.1 uniform in $s$?
clarity,_quality,_novelty_and_reproducibility: The paper is overall clear and well written. The related work is adequately discussed and the novelty of the proposed approaches both in terms of results and of techniques is clear. The paper is sufficiently original to meet the acceptance bar.

Minor clarity issues:

* Figure 1 is not particularly clear. Is there any way to highlight what is the direction of the trajectories (as $t$ grows)?

* In the discussion of the proof of Theorem 2.1 in page 6, the authors mention that "the same inequality holds with $W_0$ replaced by any $W$ s.t. $||W-W_0||\le \gamma\sqrt{m}$. What's the relationship between this $\gamma$ and $\gamma_{\rm ntk}$?

* Page 24, the last inequality of the proof in Section B.1 is incomplete.

* Lemma B.5, first inequality in the statement. $||a||\cdot ||V$ should be $||a||\cdot ||V||$. Also the authors should clarify the notation $| \mathcal X$. 


summary_of_the_review: Overall, the paper contains a number of new results fuelled by new analysis ideas. The weaknesses mentioned above are overcome by the strengths, and I am positive about the submission.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper studies the test error of the gradient flow (GF) and stochastic gradient descent (SGD) for two-layer ReLU networks. The authors provide several improved results over existing studies under several margin conditions for classification problems. Specifically, the required network width is significantly improved from $d^8$ to $d^2$ (Theorem 2.1 and 2.3) under NTK-margin assumption and from $\infty$ to $d^{d/2}$ (Theorem 3.3) under the global max margin assumption. Moreover, the lower sample complexity (Theorem 3.2) is shown under the neural collapse condition.
strength_and_weaknesses: **Strengths**:

- The improvement of the required network width (Theorem 2.1, 2.3, and 3.3) over existing results is good.
- This work develops several proof techniques: (i) tools to analyze low-width networks
near initialization, (ii) a new generalization bound technique, and (iii) a new potential function technique for global margin maximization far from initialization. These techniques will be helpful for future work.

**Weaknesses**:

The presentation is somewhat involved. 

- Although the standard margins (normalized margin and normalized smoothed margin) are introduced, they are not used in the main text. Thus, the reasons why these were introduced are unclear, which may confuse the reader.
- The explanation of the potential function on page 8 is abrupt and difficult to understand. To help readers understand the benefit of the potential function, it would be better to first explain the motivation for using it.
clarity,_quality,_novelty_and_reproducibility: The proof techniques are novel. The presentation is confusing and could be improved. 

Question: 

Does the true distribution also satisfy the neural collapse condition (Assumption 3.1) in Theorem 3.2? Assumption 3.1 appears to be described only for training dataset $(x_i,y_i)_{i=1}^n$.
summary_of_the_review: This paper makes a certain contribution in the context by improving existing results, but there is room for improvement in presentation.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: In this paper, the authors studied the optimization and generalization of two-layer ReLU network when the first layer neuron rotate little. Three regimes are considered, including the near initialization regime, Neural Collapse regime and non-rotate regime. The results show that low test error could be achieved and have some improvements over prior works in terms of either sample complexity or the width of networks. Several interesting proof ideas/techniques are discussed. Numerical experiments are also provided to support the theoretical results.
strength_and_weaknesses: Strength:
- The paper is overall well-written and easy-to-follow.
- Understanding the optimization and generalization of neural networks with finite samples is a very important problem in deep learning theory, and current paper studies it in the two-layer ReLU networks with low-rotation of the first layer and gives some interesting results.
- Some of the interesting proof ideas/techniques are presented, which uses the idea of margin. These might be of independent interest.

Weaknesses:
- While the result in Corollary 2.2 is interesting as it shows that gradient flow could escape bad KKT point in certain cases, the example here seems to be somewhat artificial in the sense that the dimension is only 2 and the labels are all +1. Given that this is a classification problem, it is a bit strange to me that the data have the same label.
- Concerns about the correctness of Theorem 2.1 and 2.3
    
    First, there is a mismatch of the initialization scale of the second layer $a$ between $a\sim N_m /\sqrt{m}$ in section 1.1 (page 2) and $a\sim N_m$ in Assumption 1.2. Under the initialization scale in Assumption 1.2, Proposition 1.3 indeed holds. Then, Proposition 1.3 was directly used in Lemma A.3. However, there seems to be an issue here: I believe Lemma A.3 in fact uses the initialization scale of $N_m/\sqrt{m}$ instead of $N_m$ because later Lemma A.3 was directly used in the proof of Theorem 2.1 and 2.3 which use the fact that $a\sim N_m/\sqrt{m}$. If the above is true, then after fixing such issue, the lower bounds in the RHS of Lemma A.3 should be scaled down by a factor of $\sqrt{m}$, which I believe would greatly change the results of Theorem 2.1 and 2.3. I would greatly appreciate if authors could comment on this to clarify my concern, and please correct me if I misunderstood any part of the proof.
- The results for the networks with large width require exponential in $d$ (input dimension) number of neurons, which makes the result not very satisfying.

Minor:
- In the line 5 of the second paragraph of page 6, $\gamma$ should be $\gamma_{ntk}$

clarity,_quality,_novelty_and_reproducibility: Clarity:
- The paper is overall clearly written and easy-to-follow.

Quality:
- The paper quality is overall good. The results are clearly stated, proof ideas are discussed, and full proofs are provided in the appendix.
- As mentioned in the weakness part, I have one concern about the proof.

Novelty:
- The theoretical results seem to be new and novel. Several interesting proof ideas/techniques are discussed.

Reproducibility:
- This is a theoretical work so there are no experiments to reproduce. Some proof ideas are discussed in the main text and full proofs are given in the appendix. I didn’t carefully check them.

summary_of_the_review: In summary, this work studied the problem of using GF and SGD on two-layer ReLU network in the classification setting and focus on the regime where the first layer neurons do not rotate much. The results show that a low test error could be achieved under different assumptions. Some of the proof ideas/techniques are also interesting. However, as mentioned in the weakness part, I currently have the concern of some part of the proof. Therefore, I’m currently leaning towards rejection. I’m willing to increase my score if authors could clarify my concerns.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: This paper studies the gradient flow/descent dynamics of a two-layer neural network with ReLU activation function and losses with an exponential tail. They derive new results for different regimes (near initialization, neural collapse, constrained layer weights that can not rotate).

The reviewers appreciate the technical novelty of the paper, although they also commented on the fact that the analysis does not provide fundamentally different results from the typical NTK analysis. Yet, there are some improvements (for instance in the terms of the requirements on the width of the network) provided by some of the theorems so this seems to be a worthwhile publication.

One reviewer identified a problem in one of the proofs which was later fixed by the authors. 

Overall, the paper definitely has some interesting results and I therefore recommend acceptance.
justification_for_why_not_higher_score: An interesting paper with somewhat novel results but not fundamentally different from prior work.
justification_for_why_not_lower_score: This is still clearly an accept, no reviewer argues for rejection
Decision: Accept: poster