Learning Hierarchical Polynomials of Multiple Nonlinear Features

We thank the reviewer for the valuable comments and suggestions.

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

Regarding the issue of assumptions, the optimality of sample complexity bound and comparison with (Nichani et al., 2023) , please see the general response for a detailed discussion.

σ1 only has a constant part and no Q1 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. We thank the reviewer for pointing out this. We will add a discussion of this equivalence between the assumption on $\sigma_1$ and adding a bias term. Specifically, if f∗=g∘pwhere p=(p1,...,pr) and each pj 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 dk, as opposed to dkq of kernel methods?

We believe that our analytic framework applies to general nonlinear features, not limited to quadratic features, such as the universality result (Lemma 1) and the proof idea proposed in Figure 1 and the entire algorithm. As the reviewer proposed, if we consider multiple degree-$q$ features and the target functions also satisfy similar assumptions as Assumptions 1-3 and the inner activation function makes emphasis on the corresponding degree ($q$) of the feature (instead of the quadratic one), our theory could be generalized to these settings with similar analysis. We just use quadratic features as an example of our theoretical framework for multiple nonlinear feature learning.

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

We thank the reviewer for mentioning the recent works on feature learning. We agree that there are emerging works beyond CSQ framework. We will add a discussion on these works in our revised version, which will come out in one or two days.

We thank the reviewer for the valuable comments and suggestions. We would appreciate it if you could champion our paper during the discussions!

Some of the assumptions might be too restrictive for the results to be relevant for practice (eg parts of Assumption 1, and Assumption 4). 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? 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. Do you expect that the d^4 dependence on dimension is necessary for quadratic features?

Regarding the issue of assumptions and the optimality of sample complexity bound (Weakness 1 and Questions 1-3), please see our general response for a detailed discussion.

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.

We agree that there exist works that analyze many-step gradient descent, but when analyzing hierarchical targets of multiple non-linear features, this could be quite challenging and requires further assumptions or modification on the network architecture like [Wang, 2023]. We believe this is a good future direction to analyze multi-step GD on the weights.

We thank the reviewer for the valuable comments and suggestions.

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.

Regarding the assumptions and comparison with previous works， please see our general response.

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 achieving small empirical loss, and 2) using Rademacher complexity analysis to bound the generalization error of . But they alone wouldn't be sufficient since we know only the existence of the 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 can be learned by gradient descent and achieves low empirical loss as well as generalization error. Even though 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

We respectfully disagree with the reviewer’s opinion that Proposition 2 is redundant for the proof of Theorem 1. To prove Theorem 1, we need to first construct an $a^{\star}$ that has both a) low norm and b) low empirical error. This implies that if we do GD on the regularized loss with appropriate choice of $\lambda$, we will converge to an $\hat a$ that has both low norm (the same scale as $a^{\star}$) and low empirical error. To conclude, we show that all the weights $a$ with low norm have a small generalization gap, which implies that the $\hat a$ we have learned has small test error. Thus, Proposition 2 is necessary for concluding Theorem 1. We will rewrite this part to make it clearer.

Minor typos

We thank the reviewer for pointing out the typos. We will fix all the typos and make notations consistent for a better readability.

There is no lower bound on the optimality of the dependence on d and κ in the bound. It should be possible to get d^2 potentially. What dependence on κ is optimal?

Regarding the optimal dependence ofn $d$, we remark that $O(d^4)$ should be optimal when considering learning our targets with one gradient step. Please see our general response for a detailed discussion. For the dependence on $\kappa$, we mainly consider balanced features cases, i.e., $\kappa$ is small and treated as nearly a constant, which is required in our analytic framework since it leads the distribution of the joint distribution of the features to resemble Gaussian distribution (Lemma 1). This allows us to compute the learned weights out approximately. When $\kappa$ is large, the distribution of the features could be quite different and this case is beyond our scope of analysis.

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.

We apologize for any confusion in the figure. We draw the figure for a visual interpretation of our proof idea presented below Proposition 1: using the universality result and the second-order information of the link function. For a better understanding of this graph, we will add a detailed explanation on how each block of terms are derived and what roles they are playing in the feature learning in our revised version. To better understand our proof idea, we recommend the readers to go over Appendix B.1, where we present an asymptotic analysis of the learned feature.

The irreducible error term of O(rp+2/d1/6) is not discussed. This is much larger than the κ1/3d1/6 term in Niching et al. (2022) and requires r to be much smaller than d Is this necessary?

A: We remark that all the previous results studying nonlinear features have an irreducible term because of the use of approximation results similar to central limit theorem on the nonlinear features to calculate the learned weights after GD. In general, the dependence on $r$ is expected and necessary when learning polynomials of $r$ features. It resembles the results of learning an arbitrary degree $p$ polynomial in $r$-dimensional space, which requires $\Omega(r^p)$ samples for efficient learning. Since $r$ represents the number of features, the case $r$ is much smaller than $d$ characterizes a typical case of feature learning, that is, the feature space is much smaller than the ambient space. if $r$ is not smaller or even bigger than $d$, the target function may be so complex that it may be only an arbitrary degree $2p$ polynomial, where feature learning is less meaningful. In this case, some generalized kernel methods like quadratic NTK [Chen, 2020] may be a more efficient analytic tool.

We appreciate the suggestions on the presentation of our work. We will emphasize the generality of our work and revise the confusing sentences for better presentation. We would be grateful if the reviewer could reconsider the evaluation based on our contributions. If you have any further questions, please do not hesitate to comment.

[Chen, 2020] Minshuo Chen, Yu Bai, Jason D. Lee, Tuo Zhao, Huan Wang, Caiming Xiong, Richard Socher Towards Understanding Hierarchical Learning: Benefits of Neural Representations

We thank the reviewer for the valuable comments and suggestions.

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.

We acknowledge that these assumptions serve to simplify technicalities, but more importantly, better illustrate the multiple feature learning process without compromising the essential nature of feature learning. Please see our general response on the assumptions. Moreover, we agree that Assumptions 2 and 3 can potentially be relaxed to some extent , such as a larger (but still near constant) tolerance factor of $\kappa_1$ and $\kappa_2$. However, we disagree that we can achieve low loss without subspace recovery. Only when we fully capture the subspace spanned by all the features, we can efficiently approximate and learn any polynomials building upon these features. Otherwise, we may suffer from a much larger sample complexity like $O(d^{2p})$ in kernel methods.

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.

We acknowledge that we do not provide a general hierarchical theoretical result like Theorem 1 in (Nichani et al. 2023), but we think this kind of result does not provide information and reveal the essence of feature learning until it is applied in specific cases to show that the components involved in the error term like $Kf^{\star}$ in (Nichani et al. 2023) can be much smaller compared with fix-feature methods. Moreover, the specific features considered in (Nichani et al. 2023) are also limited in quadratic features. Thus, in terms of end-to-end results, we do not restrict our scope. Please also see our general response for a detailed discussion.

Moreover, we remark that most of our techniques are not limited in quadratic features, such as the universality result (Lemma 1), the proof idea proposed in Figure 1 and the entire algorithm. If target functions of high-degree features also satisfy similar assumptions as Assumption 1-3 and the inner activation function makes emphasis on the corresponding degree of the feature, our theory could be generalized to these settings with similar analysis. We just use quadratic features as an example of our theoretical framework of multiple nonlinear feature learning.

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.

We thank the reviewer for pointing out the limitations of our experiments. We plan to conduct more thorough experiments to relax the assumptions and see the networks’ performance, which we will release before the end of the discussion.

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?

The uniform sphere distribution and Gaussian distribution are quite similar in high-dimensional case ($d \rightarrow \infty$), they all represent good isotropic distributions. Compared with Gaussian variables, the uniform over the sphere distribution has neater properties (some orthogonality) on kernel expansion in the $L^2$ space, as introduced in Section 2.3. Using Gaussian distribution for analysis may only differ an $o_d(1)$ term compared with the uniform sphere distribution in most of the cases. The essence of the feature learning is the same, not decided by what type of isotropic data we use.

The experimental results are still lacking, so it is hard for me to predict how they will differ from the proposed bounds.

We have conducted new experiments to verify the generality of our analytic framework under different activations, targets and relaxations of assumptions. In this link, we relax assumption 1 on how balanced the features are, i.e., the magnitude of $||A_i||$, as shown in the graph we relax the assumptions on features from fully balanced ones to extermely imbalanced ones that has spiked variance only in some directions. We see our reconstruction method in Proposition 1 actually applies to all the features, but the approximation accuracy seems to first decrease and then increase as the features get less balanced. We guess this is because the extremely imbalanced feature just reduces to multi-index model, which also has a clear sparse structure on the target.

We have also relaxed Assumption 2 from multiple quadratic features to multiple degree $q$ features and Assumption 4 from $\sigma_2=Q_q$ to $\sigma_2={\rm ReLU}$ . The results are shown in this link, where we find our method indeed applies for either ReLU activation or high-degree features.

Sample complexity dependence is sub-optimal. As the authors point out, this is perhaps unavoidable by their one-step gradient techniques, but that's not a satisfying response given that prior works have improved this in similar settings.

We understand the reviewer's criticism about the suboptimality of $O(d^4)$, but learning multiple nonlinear features with $O(d^c)$ ($c$ is a constant not dependent on $p$) was an open problem until our work. We believe our results offer sufficient contributions as the first work addressing such a broad and complex feature class. To reach $O(d^2)$ sample complexity, we may need to analyze multiple steps of GD (or SGD) to break this $O(d^4)$ gap, which is too hard to analyze.

To clarify on irreducible term ....

We remark that the irreducible error term of $O(r^{2p}/d^{1/6})$ is a result of error propagation as we have discussed under proposition 2 because the error will be amplified when we are learning degree $p$ polynomials of features rather than learning the features themselves. The error of $O(r^{p/2}/d^{1/6})$ in proposition 1 is also a result of error propagation due to the complexity of extracting the features from a degree $p$ polynomial (which is very non-smooth, $O(r^{(p-1)/2})-$Lipschitz on average) of that feature. Actually, our result in Proposition 1 completely applies to any link function $g$ rather than polynomials. Please see Proposition 3 in Appendix C.1, where we actually provide a general analysis on $L$-Lipschitz link function $g$. If the link function is smooth enough (e.g., $g$ is $O(1)-$Lipschitz) and has nondegenerate expected Hessian, the error can be reduced to $O(r^{2}/d^{1/6})$.

We would be grateful if the reviewer could reconsider the evaluation based on our contributions, our discussion on the relaxation of assumptions and our additional experiments shown above. If you have any further questions, please do not hesitate to comment.

We thank the reviewer for the valuable comments and suggestions.

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:

We understand the reviewer’s impression that our paper resembles (Nichani et al., 2023; Wang et al., 2023), but we remark that the targets of interest, neural network parametrizations, intuitions behind the results and mathematical strategies differ significantly from past works. Please see the general response part 1 for a detailed discussion on this.

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(d2), because there are O(d2) parameters to be learned in this setting. In particular, a more interesting study will consider n≍d2 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.

A:We acknowledge that $\tilde{O}(d^4)$ is suboptimal under the information theoretical pointview, but it is near optimal under our algorithm framework which only uses one gradient step for feature learning. Please see the general response part 2 for the detailed discussion of the sample complexity. Moreover, as the reviewer proposed, one can achieve tighter results like (Ba et al. 2022) in the single-index case, but our model considers multiple and nonlinear features, which are much more complex to analyze. We make the first step on analyzing this “generalized multi-index” model and provide a general analytic framework for learning multiple nonlinear features,

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.

We remark that the information-exponent is 4. We make Assumptions 2-3 to ensure a small $\mathcal{P}_{0}(f^{\star})$ and $\mathcal{P}_2(f^{\star})$ and a non-degenerate $\mathcal{P}_4(f^{\star})$. These assumptions ensure the target function adequately depends on all r features, and prevent degenerate cases, which are also considered in previous works, such as Assumption 2 in [Damian et al 2022] (where the information-exponent is 2 because they consider linear features). For the same reason, the leap-index is also basically 4 in our setting (jumping from degree 0 to degree 4).

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?

Here we want to give an example of $A_i$ that satisfies the assumption. For example, when the entries of $A_i$ follow i.i.d distributions with zero mean, right scaling of variance, and a bounded fourth moment, such as $[A_i]_{j,k}\sim N(0,1/d^2)$, then we have $||A_i||_F \approx 1$ and $||A_i||_2 \lesssim \frac{1}{\sqrt{d}}$ by standard random matrix theory.

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

The random feature embedding is commonly assumed in previous works (such as [Nichani et al., 2023]; [Wang et al., 2023], [Safran. et. al]) for multi-layer neural networks, which simplifies the analysis for the training process.

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

Assumption 4 on the activation function does not hurt the essence of feature learning, as long as the inner activation emphasizes the crucial quadratic component $Q_2$ in capturing quadratic features. Please see our general response for detailed discussions.

I thank the authors for their thorough response. While I agree (in general) with the general response part 1, I still think:

1. the setting being studied here is not novel enough.
2. there is room for improvement of the sample complexity.
3. the assumptions of the paper are unrealistic.

I will be willing to recommend the acceptance of the paper at a future time if: (1) the assumptions I discussed in my review are relaxed, or (2) the gradient update is tweaked to give a better sample complexity.

I will maintain my score for now.

We thank the reviewer for the response.

the setting being studied here is not novel enough.

We emphasize again that we consider the setting of learning multiple nonlinear features, which is not considered in any previous works (previous works either consider multiple linear feautures or single nonlinear feature). This setting also includes much broader function classes and is more interesting (and also harder) to consider. Although our algorithm framework resembles previous works, our targets of interest, neural network parametrizations, intuitions behind the results and mathematical strategies differ significantly, as we have discussed in the general response.

there is room for improvement of the sample complexity.

It is our targets of interests and algorithm framework that determine the $O(d^4)$ sample complexity. As discussed in our general response, the information exponent and leap index is basically $4$, which means that $O(d^4)$ might be neccessary for efficient recovery of the feature subspace under our algorithm framework. We understand the reviewer's criticism about the suboptimality of $O(d^4)$, but learning multiple nonlinear features with $O(d^c)$ ($c$ is a constant not dependent on $p$) was an open problem until our work. We believe our results offer sufficient contributions as the first work addressing such a broad and complex feature class. To reach $O(d^2)$ sample complexity, we may need to analyze multiple steps of GD (or SGD) to break this $O(d^4)$ gap, which is too hard to analyze.

The assumptions of the paper are unrealistic.

We have argued in the general response that Assumption 4 mentioned by the reviewer to be not realistic can actually be relaxed. We make assumption 4 just for a cleaner and centered learned weights to simplify our analysis, and the fundamental intuition is that the quadratic component $Q_2$ in the activation function plays the dominant role. We give an informal proof intuition (summarized from Appendix C.1 ). The first gradient step turn the weights into the form $h^{(1)}(x',v) \approx \sum_{i,j \ge 2}\frac{c_i^2c_j}{d^{i}} \mathbb{E}_x[f^{\star}(x) Q_i(x\cdot x') Q_j(x \cdot v) ],$ where the component $i=j=2$ will recover the quadratic features by our universality theory: $\mathbb{E}_x[f^{\star}(x) Q_2(x\cdot x') Q_2(x \cdot v) ] \approx d^{-4} \mathbf{p}(v)^{\top} \mathbf{H} \mathbf{p}(x'),$ where $\mathbf{p} \in \mathbb{R}^{r}$ is the $r$ quadratic features and $\mathbf{H}$ is the expected Hessian.

(This is also why we neeed $d^4$ samples, because we are extracting the degree $2+2=4$ component from the targets using the empirical gradients). We have proven in the Appendix C.2 and C.3 that the $(i,j)=(2,2)$ term is the dominant one under our assumptions on the targerts. Let's consider the case we relax the assumption on the activation to add back the constant and linear term ($i\le 1$ or $j \le 1$ ) .

1. The magnitude of all the terms with $i+j >4$ diminishes to $o_d(1)$ compared with the term of $(i,j)=(2,2)$
2. For the case $i+j$ is odd, the expectation is zero because $f$ has only even-degree components under our assumptions.
3. The term of $(i,j)=(0,0)$ is also $o_d(1)$ compared with the term of $(i,j)=(2,2)$ because we have asssumed (in Assumption 2 and 3) the information exponent and leap index to be basically $4$.
4. The terms of $(i,j)=(1,3)$, $(3,1)$ have the same magnitude of $(i,j)=(2,2)$, but after we apply the linear product $B^{\star}$ in Proposition, these two terms will cancel by the orthongonality of degree-$j$ term of $Q_j(x \cdot v)$ and the degree-2 terms in $B^{\star}$. This is because we simply set each row of $B^{\star}$ to be proportional $\mathbf{H}^{-1}\mathbf{p}(v)$ which only has degree $2$ terms. Thus, we still have $B^{\star} h^{(1)} (x')\approx \mathbf{H}^{-1} \mathbb{E}_{v}[\mathbf{p}(v) \cdot d^{-4} \mathbf{p}(v)^{\top} \mathbf{H} \mathbf{p}(x')] \propto \mathbf{p}(x')$
5. For $(i,j)=(2,0)$, $(4,0)$, $(0,4)$ and $(0,2)$, they introduce terms either independent of $v$ or $x'$, which leads to a common bias among the weights. (The bias of $(0,2)$ and $(2,0)$ might be even dominated by the magnitude of $(2,2)$ as discussed in case 3 ) Thus, we only need to debias the weights to remove these parts. This is also equivalent to introducing a fixed bias to the inner layer of this network as pointed out by Reviewer Xpr2.

All in all, as shown above, our assumption on inner activation does not hurt the essence of feature learning and can be relaxed. Also, these analytic strategies can naturally apply to high-degree features (not restricted to quadratic ones). We've conducted additional experiments where we use ReLU as the inner activation for comparison and we also apply our methods to high-degree features. (This link) We find that our methods indeed apply for ReLU activation and high-degree features.

2. Discussion on the sample complexity of $\tilde{O}(d^4)$

We remark that $\tilde{O}(d^4)$ is a meaningful and also reasonable sample complexity bound in our algorithm framework. Firstly, any polynomial sample complexity bound (i.e., ${O}(d^{c})$) is a significant improvement upon the bounds of $\Omega(d^{2p})$ achieved by kernel methods when learning a target of multiple quadratic features. Secondly, our $\tilde{O}(d^4)$ sample complexity is nearly optimal when considering feature learning with one gradient step. As discussed in our work, we utilize the second-order information of the link function $g$, which is dominated by the degree 4 component of the entire function, i.e. $\mathcal{P}_4(f)$. Under our assumptions on the target the “information exponent” is approximately 4 and for an efficient concentration of $\mathcal{P}_4(f)$, [Dandi et a, 2022] shows that one needs $\Omega(d^4)$ samples. This substantiates the near optimality of our results. Regarding the intuitive bound of ${O}(d^2)$ proposed by the reviewers, we think this information optimal result could be achieved when we consider different algorithms that utilize the samples more thoroughly and analyze the model asymptotically such as in the regime of $n\propto d^2$. However, this would require analyzing multiple steps of GD (or other methods) on the weights, which is quite challenging and could be a great future extension of our work.

3. Discussion on the proposed assumptions

We understand the reviewers’ concern that the four assumptions may oversimplify the model and restrict the scope of this paper, but we remark that these assumptions are all for technical simplicity, while still demonstrating the key essence of feature learning. To be specific, Assumption 1 of orthonormal quadratic features can be attained via linear transformation on the features. The assumption on the operator norm of the feature matrix represents a typical case of features, where the feature has balanced emphasis on each direction of the ambient space. The assumption on the expected Hessian in Assumption 2 and the assumption of a diminished $\mathcal{P}_2{f}$ are necessary for distinguishing all the $r$ features, which is also similarly assumed in previous works of multi-index models [Damian,2022]. A degenerate expected Hessian may fail to span all the directions spanned by the features, and a large $\mathcal{P}_2{f}$, that is, a dominant linear part of the quadratic features in the target, prevents the first step of GD on the weights from detecting the Hessian because the learned weights will concentrate to the direction of that linear part.

For the assumption on the inner activation function, we remark that the most important component is its quadratic component $Q_2$, which plays the role of capturing the quadratic features, while other components all lead to noises or biases in the learned weights. Thus, we make assumption 4 just for a cleaner and centered learned weights to simplify our analysis, the fundamental intuition is that the quadratic component plays the dominant role. Through a more complicated argument, for any inner activation function with a nondegenerate quadratic component, we can also recover the features through one gradient step on the weights.

We thank all the reviewers again for the suggestions on our work. We will add a detailed discussion on these three issues in our revised version, which will release in one or two days.

[Dandi et. al, 2022] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, Ludovic Stephan, "How Two-Layer Neural Networks Learn, One (Giant) Step at a Time"

[Damian,2022] Alexandru Damian, Jason Lee, Mahdi Soltanolkotabi, “Neural networks can learn representations with gradient descent”