Gradient-Based Feature Learning under Structured Data

Summary

Summary: The authors consider single index learning with a spiked input distribution, rather than the Gaussian typical in most works in this setting, showing that a modified SGD can learn where vanilla SGD fails and showing better sample complexity under the assumption that the spike is aligned with the true hidden signal.

Strengths

Strengths: The question of moving beyond the Gaussian input distribution is an interesting one, and the results seem quite interesting. In particular, the transitions in behavior depending on the alignment of the spike provide a rather full picture of the setting.

Weaknesses

Weaknesses: The setting is somewhat artificial, as the spike strongly correlating with the true signal is a much easier setting than the typical single index model. Furthermore, it’s not instantly clear how this version of SGD performs according to a somewhat silly baseline that’s aware of the possible correlation of the spike and true signal. From the intuition of [1], most of the effort of the learning process is spent trying to find non-negligible correlation between the learned parameter and the true direction. If you suspect the spike of your input distribution is already meaningfully correlated, you could simply initialize there. Perhaps this is a naive idea ([1] requires the link function to be known for instance), but a little discussion of why the spiked covariance setting is not too easy would strengthen the paper I feel. [1] Arous, Gerard Ben, Reza Gheissari, and Aukosh Jagannath. "Online stochastic gradient descent on non-convex losses from high-dimensional inference." The Journal of Machine Learning Research 22.1 (2021): 4788-4838.

Questions

N/A

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

N/A

Summary

In this work, the authors investigate the learning of single-index models with a gradient-based algorithm. They analyze structured inputs considering a spiked matrix model for the data. They show that an appropriate normalization of the weights in the training algorithm can help recover the target direction. They study the behavior of the sample complexity to learn the target direction as a function of the magnitude of the spike in the data, and the alignment of the spike with the target direction. They identify three learning phases, characterizing the phase transitions between them.

Strengths

The precise characterization of feature learning for single-index models is an important topic in machine learning theory. The analysis of structured input is a relevant step toward a better understanding of these models. The characterization of the sample complexity under the spiked model given in terms of a phase diagram is nice.

Weaknesses

The main weakness of the manuscript is the writing, sometimes not entirely clear.

Questions

1) The overview of the relevant literature on single-neuron training is exhaustive and nice. However, when introducing the single-neuron modification after eq. (2.3), I think it is confusing for the reader to introduce a general link $\phi$, stating assumptions on its form (e.g. Assumption 2), and then write "there is little to no hope for learning the entire single-index model unless $\phi = g$". I would rephrase this sentence in the introduction of the single-neuron simplification, given that the following results analyze non-necessarily strictly realizable settings. 2) The authors do a good job in general in trying to explain their results. I think it would be nice to expand on Proposition 4, e.g. by explaining what is the role of Assumption 2 in a sketch of proof. On a similar note, it is not clear how far from ReLU $\phi$ can be. If an intuitive understanding can be given on the class of $\phi$ to be considered, is more than welcome. 3) In Section 4.2 extends the findings to two-layer neural networks, in Thm 9 is written"suppose $\phi$ is ReLU and Assumption 2 holds", why there is a need to restrict strictly to a ReLU network in this case and not just assuming that Assumption 2 holds? 4) I like the illustrative power of Figure 1. The hard and easy phases are explained in simple terms on page 2, however, the intermediate phase is not introduced.

Rating

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

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central 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

2 fair

Contribution

3 good

Limitations

I do not see any potential negative social impact of this paper.

Summary

This paper investigates the learning dynamics of single index models when the input data has a structured covariance. It demonstrates that the commonly used spherical gradient flow fails to learn the target direction even when the spike and target directions are identical. However, the paper introduces an appropriate weight normalization technique that overcomes this limitation and successfully recovers the target direction. It further reveals that leveraging the alignment between the covariance structure and the target direction improves the sample complexity compared to isotropic cases and outperforms lower bounds for rotationally-invariant kernel methods. The paper also suggests future directions, such as studying multi-index models, understanding the limitations from a Correlational Statistical Query (CSQ) perspective, and exploring the effects of different initializations in training networks with multiple neurons.

Strengths

${\bf Originality}$: This paper introduces a new investigation into learning single-index models with a spiked covariance structure, expanding upon prior research that primarily focused on isotropic data. This unique focus on the impact of additional structure in the covariance matrix sets it apart from existing works. Furthermore, the introduction of weight normalization techniques and the exploration of their effects in anisotropic scenarios demonstrate originality in addressing the limitations of conventional spherical gradient dynamics. $\bf Quality$: This paper demonstrates a high level of quality in terms of its theoretical analysis, rigorous mathematical proofs, and well-supported findings. The authors delve into the dynamics of gradient flow, leveraging mathematical techniques to derive insights into the behavior of the learning algorithm. The use of appropriate assumptions and the incorporation of empirical covariance preconditioning showcase the paper's thoughtful methodology and attention to detail. $\bf Clarity$: The paper is written in a clear and accessible manner, making complex concepts and mathematical formulations understandable to the reader. The abstract, introduction, and conclusion concisely summarize the objectives, methods, and key findings of the study. The organization of the paper and the logical flow of ideas contribute to its overall clarity, enhancing the reader's comprehension of the research. $\bf Significance$: The paper holds good implications for the understanding of learning dynamics in the presence of structured covariance matrices. By uncovering phenomena where conventional approaches fail, such as the inability of spherical gradient dynamics to recover the target direction, the paper sheds light on the limitations and challenges associated with anisotropic data. Moreover, the proposed weight normalization techniques, insights into improved sample complexity, and the comparison to rotationally invariant kernel methods highlight the practical relevance and potential impact of the findings.

Weaknesses

${\bf (1)}\ \textbf{The training procedure considered in the paper is not practical:}$ The authors focus on a two-step training procedure that deviates from the standard gradient descent (GD) commonly used in practice. It is challenging to assess the extent to which the results obtained from this training procedure translate to impacts on standard training algorithms. ${\bf (2)} \ \textbf{Differences from the existing works:}$ Some existing works, such as Refs [BBSS22] [BAGJ21], have established some results on learning the single-index model using neural networks. It seems that the authors don't clarify the difference of this work from the existing ones. ${\bf (3)}\ \textbf{Practical implications and applications:}$ Although the authors have successfully established strong theoretical results for learning single-index models using two-layer neural networks, the practical implications of this work on real-world applications of deep learning remain somewhat unclear. Further clarification regarding the practical relevance and potential impact of the findings would greatly benefit readers seeking to understand the practical implications of the research in the field of deep learning.

Questions

$\bf Q1.$ What happens if assuming $g$ and $\phi$ to be differentiable? Can you get stronger results? $\bf Q2.$ Is it possible to extend the current analysis to the multiple-index models? What are the main technical difficulties?

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.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Not applicable

Thank you

I sincerely thank the authors for the rebuttal. After reading the authors’ response and the other reviewers’ comments, I would like to keep my score to a weak accept one.