Learning Hierarchical Polynomials with Three-Layer Neural Networks

Thank you to the reviewer for your detailed comments. We address specific concerns of yours below.

“The model and problem setting are a bit over-claim…Both problem setting and algorithm are totally far away from true three-layer neural networks or even two-layer neural networks.”

• 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.

“The results cannot handle the over-parameterized regime, i.e., the results require the width $m_1 \ll n$.”

• 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$.

“How does the first inequality in Eq. (13) hold?”

• 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.

We are happy to answer any more questions you may have, and we hope that if your concerns have been addressed you would consider raising your score.

Thank you to the reviewer for your detailed comments and for the positive evaluation of our paper. We address specific comments of yours below:

“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.

“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

Thank you to the reviewer for your detailed comments. We address specific concerns of yours below.

Please refer to the global response for discussion on the algorithm/architecture, and our new revision.

“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.”

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.

To address the questions:

• “What is the original condition on L?” It’s $L=\Theta(d)$.
• “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.
• “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].
• 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.

We are happy to answer any more questions you may have, and we hope that if your concerns have been addressed you would consider raising your score.

Amit Daniely. Depth separation for neural networks. COLT, 2017. Itay Safran, Ronen Eldan, and Ohad Shamir. Depth separations in neural networks: What is actually being separated? COLT, 2019.

Thank you to the reviewer for your detailed comments. We address specific concerns of yours below:

“The architecture and algorithm are chosen specifically to succeed for this specific class of target functions (composition of degree-k multivariate polynomials with univariate functions).”

• Please refer to the global response for discussion on the architecture/algorithm.

• We would additionally like to emphasize that a major 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 mechanism by which three-layer networks can perform feature learning. We believe that this insight is both novel and interesting.

“Several previous works have considered such layer-wise training on non-regular architectures for specific hierarchical classes of functions”

• (Allen-Zhu & Li; 2020) considers a more restricted class of neural network, in which all activations are quadratic. Furthermore, their target functions of interest are 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.
• (Abbe et al., 2020) also considers a network with quadratic activations, and a non-standard sparse architecture and training procedure that trains each neuron one-by-one. Furthermore, the “staircase” functions they consider are learnable by GD on a two-layer network with a sample/time complexity of O(d) (Abbe et al., 2022). Our work demonstrates that three-layer neural networks can learn a new class of hierarchical functions that is not known to be efficiently learnable by a two-layer network.

“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.

“For these reasons, I wonder how much this analysis can be generalized to other settings. Especially, in more realistic settings, we expect a non-linear evolution of the parameters”

• Please refer to the global comment; we would like to emphasize that prior works on two or three-layer networks also rely on layer-wise training procedures. It is straightforward to generalize our results to certain forms of non-linear parameter evolution, such as when $g_{u, V}$ is in the NTK regime.

“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$.

We are happy to answer any more questions you may have, and we hope that if your concerns have been addressed you would consider raising your score.