Summary
This paper proposes a general framework of feature learning for two-layer neural networks trained with gradient descent. This framework covers a variety of classification losses and data distributions. Specifically, the authors establish that the loss of neural networks trained with gradient descent is comparable to that of networks with first layers in the direction of a simplified gradient and optimal second layers. The framework is then specialized to multiple examples, including Gaussian mixture classification and learning parity functions.
Strengths
* The presented framework is rather general and unifies the approaches of several recent works, while covering distributions that were not previously covered in the literature.
* The studied problem is exciting and significant for the community as the existence of such a framework is necessary given the ad-hoc nature of current feature learning analyses.
* The related literature is covered extensively.
Weaknesses
* The general guarantees provided by the framework (Theorem 3.12) are limited to the performance of the optimal network with first-layer directions approximately aligned with the gradient. Such a statement does not directly imply feature learning for general problems/distributions, as a key part of feature learning seems to be showing that the gradient directions are indeed useful. To that end, the guarantees are more or less similar to prior work and suboptimal in some cases (e.g. Gaussian mixtures).
* The sample complexities of Theorems 4.4 and 4.8 are only stated as polynomials in problem parameters, thus while outperforming kernel methods, provide limited insight in comparison with prior works. Furthermore, the $\tilde{O}(d^{1.5})$ sample complexity to learn an XOR label from a mixture of 4 Gaussians in Section 4.1.1 seems to be suboptimal, as [92] shows a sample complexity of $O(d)$ for learning the same problem.
* The presentation of the paper can be improved, with specific examples given below.
Questions
* Is it possible to characterize the reason behind the suboptimality in the $\tilde{\mathcal{O}}(d^{1.5})$ sample complexity of this work in comparison with the $O(d)$ sample complexity of [92] to learn a mixture of 4 Gaussians with XOR-like structure?
* A number of figures can be added to improve the presentation of the paper, e.g. for the definition of gradient features and introducing the Gaussian mixtures and the input distribution for learning parity functions.
* Including the choices of step size in Theorem 3.12 can provide additional insights. Specifically, it seems like the theorem is still operating in the regime of a few gradient steps for the first layer, similar to prior work. A more explicit discussion of the training of the first layer would clarify the similarities and distinctions with prior work.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
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.
Limitations
The authors have mostly discussed the limitations of their work. It would help the readers better understand the limitations if there is also a discussion on how realistic the training assumption for the first layer is (i.e. only a few steps or many steps), and also the potential suboptimality in the rate of learning mixtures of Gaussians in comparison with prior work.