Summary
This paper studied the projected gradient flow on two-layer neural networks in the mean-field regime with polynomial width and quartic activation function. With data sampled uniformly on the sphere, the authors proved that to learn a single-index model with an even quartic link function, this neural network needs $n=O(d^{3.1})$ training samples. However, we know for any inner product kernel ridge regression requires $n\gg d^4$ to learn a quartic target function. This provides a concrete example of practical feature learning of neural networks outperforming the kernel methods and NTK regime when the target function has a low-rank structure.
Strengths
Overall, I found the paper well-written and easy to follow. The paper presents interesting insights into more accurate models of neural network training. The results obtained by the authors are as sharp as can be reasonably expected given the problem in certain cases. Unlike previous analyses of feature learning with two-stage training procedures, this work shows that neural networks can learn both feature direction and nonlinear link function when training in a more practical way with the projected gradient flow. The methods of the proof may provide further insights to obtain more general results on feature learning.
Weaknesses
1. The main limitation is that the activation function and the target link function are special, only quartic functions. It would be nice to check if the proof techniques can handle more general activation and target functions. Another thing that can be improved is the precise order of the width: how wide the neural network is sufficient for this feature learning.
2. There should be a comparison between the results of the current submission and [10]. [10] uses information exponent to determine the sample complexity for online SGD to learn a single-index target. When $\gamma_2=0$ in Assumption 3.2, [10] implies that sample size should be $n\gg d^{3}\log n$ which is similar to the results in the current submission. Further explanations are needed.
3. And it should be more convincing to visualize Phases 1-3 of the training dynamic in Section 5 if there is some simulation of synthetic data or real-world data for this finite-width training dynamic.
Questions
1. [Ba et al. 2022] also study the feature learning and single-index model when $n\asymp d \asymp m$ with two-stage training processes. And [Arnaboldi et al. 2023] and [Berthier et al. 2023] also studied the gradient flow dynamics of two-layer neural networks in high-dimension to a single-index model. Some comparisons with the current paper could be made.
2. In line 88, you mentioned that a similar coupling error was shown in [49] but [49] is for random feature regression with proportional limit. Can you explain more about it?
3. Line 147 typo.
4. In Theorem 3.4, the sample complexity is $n\ge d^\mu$. Is there an explicit formula or upper bound for the power $\mu$?
5. In Theorem 3.5, you considered the limitation of the inner-product kernel with data points drawn from the unit sphere. Does the data distribution have to be uniform unit sphere distribution? Can Theorem 3.5 be extended to a general rotational invariant kernel or Euclidean distance kernel?
6. Line 233 double "Appendix"
7. Line 274 $w$ should be $w_t$?
8. In the analysis, Phases 1 and 2 are controlled by the quadratic components in the activation and target functions, and the dynamics are analogous to a power method update. Is there any relation between this and the PCA warmup proposed by [Chen and Meka 2020] before SGD training?
9. In Assumption 3.2, we need $\gamma_4\ge 1.1\gamma_2^2$. Why do we need this assumption and coefficient $1.1$? What if $\gamma_2=\gamma_4=1$?
===============================================================================================
- Chen, S. and Meka, R., 2020, July. Learning polynomials in few relevant dimensions. In Conference on Learning Theory (pp. 1161-1227). PMLR.
- Ba, J., Erdogdu, M.A., Suzuki, T., Wang, Z., Wu, D. and Yang, G., 2022. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35, pp.37932-37946.
- Arnaboldi, L., Stephan, L., Krzakala, F. and Loureiro, B., 2023. From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks. arXiv preprint arXiv:2302.05882.
- Berthier, R., Montanari, A. and Zhou, K., 2023. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.00055.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
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.