Paper: Learning Hierarchical Polynomials with Three-Layer Neural Networks (ICLR-2024)

Reviewer comment (Theory):
The model and problem setting are a bit over-claim. Though it’s a three-layer neural network, only $u$ and $c$ are trained. This is actually a two-layer neural network when the input is not the original data but after a fixed feature mapping. More importantly, the used layer-by-layer training scheme makes such three-layer neural network degenerate to a composition of two kernel methods (random features). With weight decay, the loss function under two stages is strongly convex and smooth. Both problem setting and algorithm are totally far away from true three-layer neural networks or even two-layer neural networks.

Author response:
- Please refer to the global response for discussion on the architecture/algorithm. 
- We’d like to emphasize a few additional points about our work. First, even though each stage of the algorithm is a kernel method, the overall algorithm is not a kernel method, and crucially can learn hierarchical polynomials of the form $h = g \circ p$ with a much improved sample complexity. A major insight of our paper is thus that the sample complexity of learning $g \circ p$ is the same as learning $p$, by first fitting the low-degree terms of $g \circ p$ (which is approximately $p$, as proven in Lemma 2). Furthermore, this algorithm *can* be implemented by a three-layer neural network, demonstrating a mechanism by which three-layer networks *can* perform feature learning. These insights, along with the observation and also empirical evidence that a three-layer neural network can indeed implement this layerwise training procedure, has not appeared before in the literature and we believe is an interesting and important contribution.
- Next, we note that guarantees for feature learning in three-layer neural networks are few and far between, and all such results we are aware of do make some modification to the algorithm/architecture. We still find our ResNet architecture and layerwise training algorithm to be realistic while demonstrating a general feature learning principle. Our new experiments demonstrate that when all the parameters are trained jointly, the network still performs feature learning to learn the hierarchical target.

Reviewer comment (Theory):
The results cannot handle the over-parameterized regime, i.e., the results require the width $m_1 << n$. The derived results set $m_1$ and $n_1$ depending on $d$, and $m_1$ and $n$ should have different orders of $d$. Otherwise the bound will become vacuous. This is because there is one $O(\sqrt{m/n})$ term in the convergence rate, derived by Rademacher complexity for the random features model.

Author response:
- We have updated our main theorem to handle arbitrary $m_1 \ge d^k$. Since the first stage of the algorithm is a kernel method, the generalization bound should not depend on the width, and instead on the RKHS norm of the target. In the random feature setting, the RKHS norm is just the $L_2$ norm of $u^*$. Our initial proof was a bit loose, but we have refined our analysis so that by choosing the correct weight decay parameter $\xi_1$, the sample complexity now scales with the RKHS norm of $p$, independent of the width $m_1$.

Reviewer comment (Theory):
How does the first inequality in Eq. (13) hold?

Author response:
- This inequality term should actually be an equality. This is because for $j\neq l$, the two terms are independent because of the independence between $v_j$ and $v_l$, and furthermore mean zero, since $E_v[u(v)h_k(v^Tx)] = (\mathcal{P}_kh)(x)$. Therefore, the cross terms are equal to zero.

Reviewer comment (Theory):
I believe that it would be beneficial if the authors mentioned families of polynomials not captured by Assumption 4. This would make more clear how strong and restrictive this assumption is. This assumption highly simplifies the analysis and, hence, it would be nice if the authors could further discuss on this assumption (I see why the families of Remark 3 satisfy this condition, but I think a further discussion on how this assumption simplifies the analysis would be helpful).

Author response:
> “I believe that it would be beneficial if the authors mentioned families of polynomials not captured by Assumption 4… it would be nice if the authors could further discuss on this assumption”
- Here are a couple examples of polynomials which violate the assumption. One is $p(x) = x_1^3 + \cdots + x_d^3$. Even though $p$ can be written as a sum of orthogonal components, it does not lie in the span of degree 3 Hermites (as there exists a linear term). Another example is $p(x) = \sum_{i=1}^m \lambda_i h_3(x \cdot v_i)$ where $m > d$.
- Assumption 4 simplifies the analysis, in particular the proof of Lemma 2, since orthogonality of the $v_i$ implies independence of each term and thus makes it easier to compute the projection. In general, we expect that by the universality principle general features will be close to Gaussian and thus a similar result will hold, but this is challenging to make rigorous. We will be sure to add discussion on this point to a revision of our paper.

Reviewer comment (Theory):
I think a discussion on Assumption 4 (as I mention in the above section) is an important aspect that should be expanded in the current draft.

Author response:
> “I believe that it would be beneficial if the authors mentioned families of polynomials not captured by Assumption 4… it would be nice if the authors could further discuss on this assumption”
- Here are a couple examples of polynomials which violate the assumption. One is $p(x) = x_1^3 + \cdots + x_d^3$. Even though $p$ can be written as a sum of orthogonal components, it does not lie in the span of degree 3 Hermites (as there exists a linear term). Another example is $p(x) = \sum_{i=1}^m \lambda_i h_3(x \cdot v_i)$ where $m > d$.
- Assumption 4 simplifies the analysis, in particular the proof of Lemma 2, since orthogonality of the $v_i$ implies independence of each term and thus makes it easier to compute the projection. In general, we expect that by the universality principle general features will be close to Gaussian and thus a similar result will hold, but this is challenging to make rigorous. We will be sure to add discussion on this point to a revision of our paper.

Reviewer comment (Theory):
The current result heavily relies on Gaussian marginals. Do you think similar results could be established for other continuous measures? Or discrete ones (in the Boolean domain, any function is essentially a polynomial in the Fourier basis)?

Author response:
> “Do you think similar results could be established for other continuous measures?”
- We believe the high level insights should be the same; however, we note that in the literature essentially all feature learning results require strong assumptions on the data distribution such as being Gaussian, uniform on the hypersphere, or uniform on the Boolean hypercube. We expect our results to translate to these simple distributions, but for more general distributions this question is still open for much simpler models such as single-index models (see for example (Bruna et al., 2023) for preliminary work in this direction)

Joan Bruna, Loucas Pillaud-Vivien, Aaron Zweig. On Single Index Models beyond Gaussian Data. NeurIPS 2023. arxiv.org/abs/2307.15804

Reviewer comment (Novelty):
The main weakness of the paper is the close relationship with previous works. Although a fair comparison is given, the works by [Allen-Zhu & Li, 2019;2020] and [Nichani et al. 2023] contain many of the key ideas in the manuscript.

Author response:
We would like to emphasize that the major original insight of our paper is that the sample complexity of learning $g \circ p$ is the same as learning $p$, by first fitting the low-degree terms of $g \circ p$ (which is approximately $p$, as proven in Lemma 2). Furthermore, this algorithm can be implemented by a three-layer neural network, demonstrating a new mechanism by which three-layer networks can perform feature learning. We believe that this insight is both novel and interesting.

To compare to the specific prior works:
- (Allen-Zhu & Li; 2019; 2020) considers a more restricted class of neural network, in which all activations are quadratic. Furthermore, they consider target functions of the form $h(x) = p + \alpha g \circ p$, where $\alpha = o_d(1)$. As such, the learned feature $p$ already achieves vanishing test loss. Our setting is much more challenging, in that the learner must extract $p$ while only having access to $h = g \circ p$; the learner must also fit $g$ in order to obtain vanishing test error. We emphasize here that our choice of link function $g$ is very general, while Allen-Zhu & Li; 2020 can only handle $g$ which are almost linear. Our insight that a three-layer neural network can indeed implement both of these steps while only having access to $h$ is indeed a novel contribution compared to this prior work. 
- Likewise, (Nichani et al., 2023) can only handle quadratic features, and for that setting obtains a suboptimal sample complexity ($d^4$). Our key contribution is that this general principle of hierarchical feature learning -- the sample complexity of learning $g \circ p$ via a three-layer neural network is the same as learning $p$ -- holds for a broad class of hierarchical functions. A major technical innovation of our work is an extension of the “approximate Stein’s lemma” argument to handle degree $k$ polynomials.

Reviewer comment (Theory):
In remark 2 is written "can be extended to any $L= \omega_d(1)$". What is the original condition on L?

Author response:
“What is the original condition on L?” It’s $L=\Theta(d)$.

Reviewer comment (Theory):
Could the author comment with references characterizing the failure of two-layer networks in learning these high-rank target functions? Are there any provable guarantees?

Author response:
“Could the author comment with references characterizing the failure of two-layer networks in learning these high-rank target functions?” Existing guarantees for two-layer neural networks fall in two rough categories. One set of results uses NTK style arguments; while these can handle high-rank target functions, their sample complexity is no better than a kernel method. On the other hand, feature learning guarantees like [Damian et al., 2022; Abbe et al., 2022; 2023] which do improve over kernel methods require the target to be low rank. While no general lower bound against two-layer networks learning high rank functions exists, approximation lower bounds have been proven against specific high-rank targets: see for instance [Daniely, 2017; Safran et al. 2019; Safran & Lee, 2022; Nichani et al., 2023]. We will update the exposition to make this point more clear.

Reviewer comment (Theory):
Before Lemma 1 mention that $n = O(d^k)$ is used, and repeat references that prove for what reason you fit the best degree k polynomial.

Author response:
We will be sure to add additional discussion on [Ben Arous et al., 2021], compare to two-layer networks after Corollary 1, and mention the $n = d^k$ requirement before Lemma 1.

Reviewer comment (Theory):
Could the author comment on the requirements for the gradient step size in the first phase? It is mentioned that [Nichani et al. 2023] considered one large gradient step in this phase. How does this relate to your learning rate requirement in Thm 1?

Author response:
“Could the author comment on the requirements for the gradient step size in the first phase?” We need the learning rate to be small enough to mimic gradient flow -- $\eta < 1/\ell$, where $\ell $ is the smoothness of the loss. This is a different mechanism to [Nichani et al., 2023], which requires a very large gradient step that scales inversely to the initialization size. We believe our algorithm better captures practice, where smaller gradient steps to mimic gradient flow are more common. This distinction is another point of novelty of our work as compared to [Nichani et al., 2023].

Reviewer comment (Novelty):
Overall, the architecture choice and algorithm makes the analysis straightforward. It reduces the problem to sequentially fitting two linear random feature models, which is by now quite well understood. See for example [Generalization error of random feature and kernel methods: Hypercontractivity and kernel matrix concentration, Mei et al., 2022] which gives the sample size and network width $\Theta(d^k)$ to learn any degree-$k$ polynomials (for spherical data). The main technical innovation is Lemma 2, which relies on the specific construction of assumption 4.

Author response:
> “Overall, the architecture choice and algorithm makes the analysis straightforward. It reduces the problem to sequentially fitting two linear random feature models, which is by now quite well understood.”

- While each stage of our algorithm indeed fits a random feature model, the overall algorithm of sequentially fitting random feature models is not a kernel method, and crucially can learn hierarchical polynomials of the form $h = g \circ p$ with a much improved sample complexity. This fact, along with the observation that a three-layer neural network can indeed implement this layerwise training procedure, has not appeared before in the literature and we believe is an interesting and important contribution.

Reviewer comment (Theory):
Why not directly consider a general function $g$ instead of a polynomial? The random feature analysis shouldn’t change much in that case (it simply requires to show a $\eps$-approximate certificate using uni-dimensional random features).

Author response:
> “Why not directly consider a general function g instead of a polynomial? The random feature analysis shouldn’t change much in that case.” 

- Currently, our proof of Lemma 2 requires expanding $g \circ p$ into a sum of products of monomial terms, and computing the projection of each individual term. As such, the current proof only applies for polynomials. However, we expect the lemma to still hold for more general $g$, albeit with some more technical challenges to get the proof to go through. Overall, the high level insight that three-layer neural networks perform hierarchical feature learning should still remain for general $g$.