Learning Hierarchical Polynomials of Multiple Nonlinear Features

Authors: Hengyu Fu, Zihao Wang, Eshaan Nichani, Jason D. Lee
Primary area: learning theory
Keywords: Deep Learning Theory, Feature Learning, Three-Layer Neural Network, Gradient Descent, Sample Complexity

Abstract:
In deep learning theory, a critical question is to understand how neural networks learn hierarchical features. In this work, we study the learning of hierarchical polynomials of multiple nonlinear features using three-layer neural networks. We examine a broad class of functions of the form $f^{\star}=g^{\star}\circ \mathbf{p}$, where $\mathbf{p}:\mathbb{R}^{d} \rightarrow \mathbb{R}^{r}$ represents multiple quadratic features with $r \ll d$ and $g^{\star}:\mathbb{R}^{r}\rightarrow \mathbb{R}$ is a polynomial of degree $p$. This can be viewed as a nonlinear generalization of the multi-index model, and also an expansion upon previous work on nonlinear feature learning that focused only on a single feature (i.e. $r = 1$). 
    Our primary contribution shows that a three-layer neural network trained via layerwise gradient descent suffices for  
 - complete recovery of the space spanned by the nonlinear features 
- efficient learning of the target function $f^{\star}=g^{\star}\circ \mathbf{p}$ or transfer learning of $f=g\circ  \mathbf{p}$ with a different link function
   
within $\widetilde{\mathcal{O}}(d^4)$ samples and polynomial time.
For such hierarchical targets, our result substantially improves the sample complexity ${\Theta}(d^{2p})$ of the kernel methods, demonstrating the power of efficient feature learning. It is important to highlight that our results leverage novel techniques and thus manage to go beyond all prior settings such as single-index and multi-index models as well as models depending just on one nonlinear feature, contributing to a more comprehensive understanding of feature learning in deep learning.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1730691727108

thread_type:
official_review

<thread_object>
creator: reviewer_xpr2

<content>
summary:
This work studies the problem of learning functions $f^* : \mathbb{R}^d \to \mathbb{R}$ of the form $f^* = g^* \circ p$ where $p : \mathbb{R}^d \to \mathbb{R}^r$ returns $r$ quadratic features of the input, and $g^* : \mathbb{R}^r \to \mathbb{R}$ is a polynomial of degree $p$. By analyzing a gradient-based algorithm for training a three-layer neural network with specific choices of activation, the authors prove that a sample complexity of $\tilde{O}(d^4)$ is sufficient for learning this function, as opposed to the $\Omega(d^{2p})$ sample complexity lower bound for learning with rotationally invariant kernels.

strengths:
Presenting a guarantee for learning a polynomial of multiple quadratic features using three-layer neural networks is a new direction in the literature which extends beyond the guarantees for learning multi-index models with two-layer neural networks and learning polynomials of a single non-linear feature which existed in prior works. More broadly, this work contributes to finding settings in which networks deeper than two layers are natural for learning certain function classes, which is an area of particular interest for the community.

weaknesses:
* The feature learning regime considered in this paper is limited in the sense that the first layer weights are fixed at random and not trained. Training these weights could lead to learning more complex hierarchies, such as hierarchies of multi-index models, although admittedly challenging for a theoretical analysis.

* The choice of activation for the second layer is rather specific, and choosing a different activation for each layer seems a bit odd when compared with standard three-layer neural networks. However, I believe there might be a way to choose the same activation as my impression is that $\sigma_1$ only has a constant part and no $Q_1$ part, thus one only needs to introduce a fixed bias to the output of the network. If this interpretation is correct this might be worth mentioning in the main text.

* While the statistical model is naturally suitable for three-layer networks, I could not find any lower bounds in this submission which explain that learning this function with a two-layer neural network is provably less efficient (in terms of samples or computation). I am wondering if the authors could use the results of prior works to show a statistical or computational lower bound holding in this setting.

questions:
* What is the main technical challenge in directly generalizing the result of Nichani et al., 2023 to the setting of having multiple non-linear features? Specifically, if $f^* = g \circ p$ where $p = (p_1,...,p_r)$ and each $p_j$ is a degree $k$ polynomial of $x$, and $g$ is a degree $q$ polynomial, is it possible to argue that in the regime of constant $r$, the sample complexity of learning $f^*$ with three-layer neural networks is $d^k$, as opposed to $d^{kq}$ of kernel methods?

* On the discussion of learning linear features, a relevant work may be [1] that considers learning general multi-index models with two-layer neural networks via a noisy gradient descent algorithm. For single-index models, while the authors have included works in the CSQ framework, there are more recent advances on gradient-based learning with two-layer neural networks that go beyond the CSQ framework [2, 3], or consider structured rather than isotropic input covariance [4, 5].

---
[1] A. Mousavi-Hosseini et al. "Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics". arXiv 2024.

[2] Y. Dandi et al. "The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents". ICML 2024.

[3] J. D. Lee et al. "Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit". arXiv 2024.

[4] A. Mousavi-Hosseini et al. "Gradient-Based Feature Learning under Structured Data". NeurIPS 2023.

[5] J. Ba et al. "Learning in the presence of low-dimensional structure: a spiked random matrix perspective". NeurIPS 2023
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699113336

thread_type:
official_review

<thread_object>
creator: reviewer_bxau

<content>
summary:
The authors prove that three-layer networks can learn nonlinear, multi-index targets more efficiently than kernel methods. Quite a bit of literature already exists on 2- and 3-layer networks learning single-index targets of linear and non-linear features, and also multi-index targets of linear features. To the best of my knowledge this is the first set of results for learning multiple, non-linear features. 

The analysis is valid of an algorithm that does a single step on the first layer weights, W, and then completely trains the rest of the network. The authors use this two-stage training methodology to also prove a transfer learning result, where the link function is changed. 

Finally the authors provide some empirical evidence that supports their theory. In particular, their approach appears to improve sample complexity over a random feature model. They also support their findings regarding efficient transfer learning and the reconstruction of quadratic features.

strengths:
This is a strong submission in most respects: relevance, rigor, and clarity. I have minor comments and questions below.

weaknesses:
1. Some of the assumptions might be too restrictive for the results to be relevant for practice (eg parts of Assumption 1, and Assumption 4). 
2. Algorithm 1 does but a single update on W. Then parameter a is trained on multiple steps. While this style of analysis is common in the area, there already exist articles that include analysis of many-step gradient descent for all parameters.

questions:
1. The “balanced features” part of Assumption 1, where the operator norm of Ai satisfies a particular bound. This bound essentially tells us that each feature cannot depend on a sparse subset of inputs. Can you discuss the implications of that assumption and pertinence to real distributions? 
2. Similarly, can you discuss the implications of Assumption 4, especially the Gegenbauer expansion? Activation functions used in practice are not even, and many of them are actually odd.
3. Do you expect that the d^4 dependence on dimension is necessary for quadratic features?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730706793345

thread_type:
official_review

<thread_object>
creator: reviewer_okdb

<content>
summary:
In this work, the authors studies theoretically the setting of using three-layer neural networks to learn a type of "multi-index quadratic models" under a uniform data distribution on the sphere. With suitable activation functions and layer-wise training, the model is proved to be able to learn both the output function and the quadratic features with a sample complexity that scales as $d^4$.

strengths:
The type of model studied in this paper generalizes the one considered by Nichani et al. (2023) by allowing multiple quadratic features instead of just one in the first layer. To prove the learning guarantee, the authors introduced a nice technical result on the distribution of multiple nonlinear functions of Gaussian inputs, which is interesting mathematically and novel to my knowledge. The theoretical results appear to be sound though they have not been thoroughly checked.

weaknesses:
**Beyond the mathematical and technical**: It could be argued that the series of assumptions on the layer-wise learning algorithm, the activation functions, the data distribution, etc. are quite specific and restrictive, and the contribution of this work (relative to e.g. Nichani et al., 2023 and Wang et al., 2023, which share much of a common analytical framework) resides more in the improved mathematical results than advancements in new learning theory insights.

**Some confusion regarding the proof of Theorem 1**: The authors' description of the second stage in the proof of Theorem 1 gives the impression that it breaks down into two parts: 1) showing the existence of an $a^*$ achieving small empirical loss, and 2) using Rademacher complexity analysis to bound the generalization error of $a^*$. But they alone wouldn't be sufficient since we know only the existence of the $a^*$ but not its learnability by the proposed algorithm. And indeed, the actual proof steps taken by the authors is different (and correctly so), where they use the strong convexity of the loss function to show that an alternative solution $\hat{a}$ can be learned by gradient descent and achieves low empirical loss as well as generalization error. Even though $a^*$ still makes its appearance in the proof argument in Section D.1.2, I suspect that it is not truly necessary and only complicates the proof. If this assessment is correct, it would imply that Proposition 2 is redundant for the proof of Theorem 1, which should be commented on in order to avoid confusion.

Minor typos:
1) In Section 1.1 on Page 2, "outer layer can achieves" --> "outer layer can achieve".
2) In Section 4.1 on Page 7, "trained presentations representations" --> "trained representations".
3) The links referencing Proposition 6 are misplaced.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730726122457

thread_type:
official_review

<thread_object>
creator: reviewer_buve

<content>
summary:
The paper studies hierarchical feature learning in fully-connected neural networks when trained with gradient descent. In particular, the authors consider a nonlinear (quadratic) multi-index model with $f^*(x) = g^*(x^\top A_1 x, \ldots, x^\top A_r x)$ for $x \in \mathcal{S}^{d-1}$ and $r\ll d$. Under approximate orthogonality, well-balancedness conditions of the quadratic features, and additional well-conditioned assumptions on the link function, the authors show that a layer-wise gradient descent method on a three-layer neural network (first layer is never trained, and activations are carefully chosen) can learn $f^*$ with roughly $\tilde{O}(d^4 + d^2r^{p+1})$ (ignoring $\kappa$) to irreducible error $O(r^{p+2}/d^{1/6})$, and approximately recover the $r$-dimensional subspace to allow for transfer learning to a different link function.

strengths:
- The paper studies an interesting model of feature learning, generalizing prior work that consider either linear multi-index settings or non-linear single-index models.
- The authors show that gradient descent recovers the subspace which allows them to give transfer learning results for any link function.
- To prove the results, the authors prove a new universality theorem for multiple-dimensions, generalizing previous known results. Though the proof seems to follow from existing tools in high dimensional probability, it seems novel.

weaknesses:
- The paper makes several assumptions and tweaks that seem to be done to make the proofs easier: layer-wise training, removing lower order terms in activation functions, taking almost orthonormal equally important features, etc. Most of these are standard is deep learning theory, but the activation assumption seems more artificial, and tailored to make the analysis easier. Some assumptions (Assumptions 2 and 3) can potentially be relaxed if the goal is only getting low loss and not subspace recovery.
- Compared to prior work (Nichani et al. 2022) that was for general hierarchal functions, this paper is narrow in its scope and relies heavily on tools already available in this prior work, while requiring additional assumptions. I do understand that the prior work did not look at multiple quadratic features, but limiting to quadratic seems like a limitation of the techniques.
- The experiments are limited to following the theoretical assumptions precisely. This is not very insightful. If the experiments are to remain synthetic, then it would be worthwhile seeing how the scaling with different parameters in the bounds shows up in experiments. Also, how much can you relax the assumptions, and still get the same empirical results. It would also be worthwhile to come up with a real-world dataset where this quadratic model makes sense, and gains of three vs two layers can be shown.

questions:
- Why did the authors consider only uniform over the sphere compared to Gaussian, which is more commonly studied in this literature? Do the results extend to more general distributions?
- There is no lower bound on the optimality of the dependence on $d$ and $\kappa$ in the bound. It should be possible to get $d^2$ potentially. What dependence on $\kappa$ is optimal?
- Figure 1 is not very easy to understand. Might be helpful to add more explanation to the caption and label the different blocks/terms. It took me a while to parse it and I did not find it very helpful to understand the proof. 
- The irreducible error term of $O(r^{p+2}/d^{1/6})$ is not discussed. This is much larger than the $\frac{\kappa^{1/3}}{d^{1/6}}$ term in Niching et al. (2022) and requires $r$ to be much smaller than $d$. Is this necessary?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731267398340

thread_type:
official_review

<thread_object>
creator: reviewer_jyeu

<content>
summary:
The authors study layerwise training on a three layer neural network trained on a more general target function with a more involved "hierarchical feature" structure. They show sample complexity improvements over the classic kernel methods.

strengths:
- Feature learning is one of the most fundamental problems in deep learning theory.

- The proof seem to be correct. Nice proof sketches are provided in the main text.

weaknesses:
- I don't find the results of this paper particularly interesting. I would have easily guessed that these results should be obtainable for the case where r \neq 1, after looking at Nichani et al., 2023; Wang et al., 2023. See also my other comments below:

- The results of the paper are not tight; i.e., the sample complexity of O(d^4). Essentially, one would expect a sample complexity of $O(d^2)$, because there are $O(d^2)$ parameters to be learned in this setting. In particular, a more interesting study will consider $n \asymp d^2$ and provide a precise characterization of the limiting test error. In the single index case, this is similar to the benefit show by Ba et al. 2022 (and others) over the initial results of Damian et al 2022.

- In assumption 2, it is assumed that the Hessian has eigenvalues that are bounded below by a certain quantity. Can this be translated into a condition on the target function in terms of some quantity like information-exponent or leap-index? Right now, the assumption seems very artificial.

- What do the authors mean by " Moreover, when the entries of Ai are sampled i.i.d., the assumption is satisfied with high probability by Wigner’s semicircle law" in line 157?

- In the layer-wise training, the inner most layer is not trained and is kept at initialization; i.e., random features.

- Assumption 4 on the activation function is completely unrealistic and is only an artifact of the proof.

- What is the role of the update on the bias. Some prior work in the area consider neural networks without bias (e.g., the papers I list below). Can the authors comment on this please?

- The discussion of prior work on "Learning Linear Features" is very limited. There are many more results on this topic, studying feature in this case with detail. The authors should consider citing and discussing these papers (alongside any other paper on the topic that I might have missed here).

 [1] Dandi et al., The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents

 [2] Moniri et al., A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks

 [3] Cui et al., Asymptotics of feature learning in two-layer networks after one gradient-step

 [4] Dandi et al., A Random Matrix Theory Perspective on the Spectrum of Learned Features and Asymptotic Generalization Capabilities

 [5] Wang et al, Nonlinear spiked covariance matrices and signal propagation in deep neural networks

questions:
please see the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732259769808

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to All Reviewers Part 1

comment:
We thank all reviewers for the detailed comments and feedback. We’d first like to address comments that were shared by a few of the reviewers.

**1.Comparison with  (Nichani et al., 2023) & (Wang et al., 2023)**

While our work seems similar to these two previous works in the algorithm framework, we point our that the **targets of interests**, the **parametrizations of neural networks**, **intuitions behind the results** and the corresponding **mathematical strategies** are of great difference compared to the past works, not an extension to either (Nichani et al., 2023) or (Wang et al., 2023). The theory of (Nichani et al., 2023) or (Wang et al., 2023) cannot be generalized to our work with an “easy guess”.

**a) Target of interests**: We consider learning hierarchical polynomials
 of multiple quadratic features, while  (Nichani et al., 2023) considers polynomials of single quadratic feature and   (Wang et al., 2023) consider broader feature class but still limited to a  single feature. Our function class contains  **multiple** non-linear features, which have not been considered in any previous work. This generalized multi-index target is much more interesting and also **inherently harder** to analyze compared with these two works. Moreover, our assumptions on the target functions, particularly the link function $g$ are quite different and even contrary to the previous two works. While  (Nichani et al., 2023) & (Wang et al., 2023) assumes the link function $g$ has non-zero linear component, we rather assume in Assumption 2 & 3 $g$ has a **non-degenerate expected Hessian** and a nearly **zero linear component**. This is the first **essential difference between single-feature and multi-feature** learning, because  our assumptions ensure that the network genuinely learns to represent and distinguish all $r$ features, while assumptions proposed by  (Nichani et al., 2023) & (Wang et al., 2023)  represent a degenerate case that neural network may only learn the dominant linear combinations of the  $r$ features. We’ve also provided examples and counterexamples in Remark 1 to elaborate on this difference.


**b) Parametrizations of neural networks**: Different from (Nichani et al., 2023) and (Wang et al., 2023), both which initialize the second-layer weights (i.e., $\mathbf{W}$ in our paper) at zeros, we initialize $\mathbf{W}$ as **independent Gaussians**. This **random initialization** instead of a deterministic one is another essential difference between multiple feature learning and single feature learning, because we require the weights to **span the full space**, which enables the learned weights (after one step of GD) to **capture the multiple features in all directions** instead of converging to a specific direction like (Nichani et al., 2023) and (Wang et al., 2023), i.e., single feature learning.

**c) Intuition and mathematical strategies**: The intuition of (Nichani et al., 2023) and (Wang et al., 2023) is that the gradient descent on the zeroed weights $W$ of  (either one step or multiple steps) make them to concentrate to the single target feature (scaled by a constant). On the contrary under our assumptions on the targets and proper initialization on the weights of neural network，our intuition is that with sufficient number of training samples, gradient step on the randomly initialized weights $W$, i.e., Stage 1,  makes them **fully recover and span the $r-$dimensional space** of all the non-linear features, as discussed in Proposition 1 and Figure 1. **Subspace recovery is completely different from and also significantly harder than single feature recovery**.

 
To prove this statement, we propose a **novel universality theory** (Lemma 1) that relates the joint distribution of multiple nonlinear features to a multivariate Gaussian, which enables us to approximately compute the learned weights. This technique of multivariate Gaussian approximation is not considered in any previous work. More specifically, to handle learning multiple features, the proof of Proposition 1 is different and also more complex than similar results in (Nichani et al., 2023) and (Wang et al., 2023) for learning a single-feature, because we are focusing on the calculation of second-order terms of the link function, i.e. $\mathbb{E}]\nabla^2 g(p)]$, while their works mainly consider the first-order term  $\mathbb{E}]\nabla g(p)]$. Our computation requires more techniques in high dimensional probability and random matrix theory.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732538236779

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Revision of our paper comes out

comment:
Dear reviewers,

Thank you again to all the reviewers for your detailed feedback on our submission. We have uploaded the revised version of our paper, in which we have added detailed **comparison with previous works** (Nichani et al., 2023) & (Wang et al., 2023), **more recent related works** and **explanations on our proposed assumptions**. Due to the page limit, we have moved the original Section 5 (Numerical Experiments) to Appendix A. All the modifications in the main page are **marked in red** for the reviewers’ convenience.

Given that the discussion period ends soon, we would like to know whether our rebuttal has addressed your concerns, and whether any more questions of yours remain. Please do let us know, we are happy to answer any additional questions.

Best, 

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper considers the problem of learning of hierarchical polynomials of multiple nonlinear features using three-layer neural networks. Authors consider a layer-wise training procedure and prove recovery of the space spanned by the nonlinear features, and learning the target non-linearity on top. The overall sample complexity of the procedure is d^4 with polynomial-time iteration.


This paper was reviewed by five expert reviewers with the following Scores/Confidence: 8/4, 6/4, 6/4, 8/4, 6/2. I think the paper is studying an interesting topic and the results are relevant to ICLR community. The following concerns were brought up by the reviewers:

- The algorithm in this paper deviates from deep learning practice in many parts. Although the methodology is clearly explained, these should be emphasized in the introduction and throughout the main text.
 
  
- While I agree with the authors that d^4 sample complexity is reasonable, I suggest they include discussions on any room for improvements. This is a point that was raised by multiple reviewers and reader could also benefit.
  
- several typos and ambiguous statements were pointed by the reviewers. These should be carefully addressed.


Authors should carefully go over reviewers' suggestions and address any remaining concerns in their final revision. Based on the reviewers' suggestion, as well as my own assessment of the paper, I recommend including this paper as 'spotlight' to the ICLR 2025 program.