SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} f \left(x ^ {\prime}; \theta^ {(1)}\right) = \frac {1}{m _ {1}} \sum_ {j = 1} ^ {m _ {1}} a _ {j} \sigma_ {1} \left(\left\langle w _ {j} ^ {(1)}, h ^ {(0)} \left(x ^ {\prime}\right) \right\rangle + b _ {j}\right) (8) \\ = \frac {1}{m _ {1}} \sum_ {j = 1} ^ {m _ {1}} a _ {j} \sigma_ {1} \left(a _ {j} ^ {(0)} \cdot \bar {\eta} \frac {1}{m _ {2}} \mathbb {E} _ {x} \left[ f ^ {*} (x) h ^ {(0)} (x) ^ {T} h ^ {(0)} \left(x ^ {\prime}\right) \right] + b _ {j}\right) \mathbf {1} _ {b _ {j} ^ {(0)} \geq 0}. (9) \\ \end{array}
$$
We first notice that this network now implements a 1D function of the quantity
$$
\phi \left(x ^ {\prime}\right) := \bar {\eta} \frac {1}{m _ {2}} \mathbb {E} _ {x} \left[ f ^ {*} (x) h ^ {(0)} (x) ^ {T} h ^ {(0)} \left(x ^ {\prime}\right) \right]. \tag {10}
$$
Specifically, the network can be rewritten as
$$
f \left(x ^ {\prime}; \theta^ {(1)}\right) = \frac {1}{m _ {1}} \sum_ {j = 1} ^ {m _ {1}} a _ {j} \sigma_ {1} \left(a _ {j} ^ {(0)} \cdot \phi \left(x ^ {\prime}\right) + b _ {j}\right) \cdot \mathbf {1} _ {b _ {j} ^ {(0)} \geq 0}. \tag {11}
$$
Since  $f$  implements a hierarchical function of the quantity  $\phi(x)$ , we term  $\phi$  the learned feature.
The second stage of Algorithm 1 is equivalent to random feature regression. We first use results on ReLU random features to show that any  $q \in \mathcal{W}^{2,\infty}([-1,1])$  can be approximated on  $[-1,1]$  as  $q(z) \approx \frac{1}{m_1} \sum_{j=1}^{m_1} a_j^* \sigma_1(a_j^{(0)}z + b_j)$  for some  $\|a^*\| \lesssim \|q\|_{2,\infty}$  (Lemma 3). Next, we use the standard kernel Rademacher bound to show that the excess risk scales with the smoothness of  $q$  (Lemma 5). Hence we can efficiently learn functions of the form  $q \circ \phi$ .
It suffices to compute this learned feature  $\phi$ . For  $m_2$  large, we observe that
$$
\phi \left(x ^ {\prime}\right) = \frac {\bar {\eta}}{m _ {2}} \sum_ {j = 1} ^ {m _ {2}} \mathbb {E} _ {x} \left[ f ^ {*} (x) \sigma (x \cdot v) \sigma \left(x ^ {\prime} \cdot v\right) \right] \approx \bar {\eta} \mathbb {E} _ {x} \left[ f ^ {*} (x) K \left(x, x ^ {\prime}\right) \right] = \bar {\eta} (\mathbb {K} f ^ {*}) \left(x ^ {\prime}\right). \tag {12}
$$
The learned feature is thus approximately  $\overline{\eta} \cdot \mathbb{K}f^{*}$ . Choosing  $\overline{\eta}$  so that  $|\phi(x')| \leq 1$ , we see that Algorithm 1 learns functions of the form  $\phi \circ (\overline{\eta} \cdot \mathbb{K}f^{*})$ . Finally, we translate the above analysis to the finite sample gradient via standard concentration tools. Since the empirical estimate to  $\mathbb{K}f^{*}$  concentrates at a  $1/\sqrt{n}$  rate,  $n \gtrsim \| \mathbb{K}f^{*}\|_{L^2}^{-2}$  samples are needed to obtain a constant factor approximation (Lemma 7).
We next instantiate Theorem 1 in two specific statistical learning settings which satisfy the hierarchical prescription detailed in Section 3. As a warmup, we show that three-layer networks efficiently learn single index models. Our second example shows how three-layer networks can obtain a sample complexity improvement over existing guarantees for two-layer networks.
4.1 Warmup: single index models
Let  $f^{*} = g^{*}(w^{*}\cdot x)$ , for unknown direction  $w^{*}\in \mathbb{R}^{d}$  and unknown link function  $g^{*}:\mathbb{R}\to \mathbb{R}$ , and take  $\mathcal{X}_d = \mathbb{R}^d$  with  $\nu = \mathcal{N}(0,I)$ . Prior work [18, 13] shows that two-layer neural networks learn such functions with an improved sample complexity over kernel methods. Let  $\sigma_{2}(z) = z$ , so that the network is of the form  $f(x;\theta) = \frac{1}{m_1} a^T\sigma_1(WVx)$ . We can verify that Assumptions 1 to 4 are satisfied, and thus applying Theorem 1 in this setting yields the following:
Theorem 2. Let  $f^{*}(x) = g^{*}(w^{*}\cdot x)$ , where  $\| w^{*}\|_{2} = 1$ . Assume that  $g^{*},(g^{*})^{\prime}$  and  $(g^{*})^{\prime \prime}$  are polynomially bounded and that  $\mathbb{E}_{z\sim \mathcal{N}(0,1)}[(g^{*})'(z)]\neq 0$ . Then with high probability Algorithm 1 satisfies the population loss bound
$$
\mathbb {E} _ {x} \left[ \left(f (x; \hat {\theta}) - f ^ {*} (x)\right) ^ {2} \right] = \tilde {O} \left(\frac {d ^ {2}}{\operatorname* {m i n} \left(n , m _ {1} , m _ {2}\right)} + \frac {1}{\sqrt {n}}\right). \tag {13}
$$
Given widths  $m_{1}, m_{2} = \tilde{\Theta}(d^{2})$ ,  $n = \tilde{\Theta}(d^{2})$  samples suffice to learn  $f^{*}$ , which matches existing guarantees for two-layer neural networks [13, 18]. We remark that prior work on learning single-index models under assumptions on the link function such as monotonicity or the condition  $\mathbb{E}_{z \sim \mathcal{N}(0,1)}[(g^{*})'(z)] \neq 0$  require  $d$  samples [49, 37, 12]. However, our sample complexity improves on that of kernel methods, which require  $d^{p}$  samples when  $g^{*}$  is a degree  $p$  polynomial.
Theorem 2 is proved in Appendix E.1; a brief sketch is as follows. Since  $\sigma_2(z) = z$ , the kernel is  $K(x, x') = \mathbb{E}_v[(x \cdot v)(x' \cdot v)] = \frac{x \cdot x'}{d}$ . By an application of Stein's Lemma, the learned feature is
$$
\left(\mathbb {K} f ^ {*}\right) (x) = \frac {1}{d} \mathbb {E} _ {x ^ {\prime}} \left[ x \cdot x ^ {\prime} f ^ {*} (x) \right] = \frac {1}{d} x ^ {T} \mathbb {E} _ {x ^ {\prime}} \left[ \nabla f ^ {*} \left(x ^ {\prime}\right) \right] \tag {14}
$$
Since  $f^{*}(x) = g^{*}(w^{*}\cdot x)$ ,  $\nabla f^{*}(x) = w^{*}(g^{*})^{\prime}(w^{*}\cdot x)$ , and thus
$$
\left(\mathbb {K} f ^ {*}\right) (x) = \frac {1}{d} \mathbb {E} _ {z \sim \mathcal {N} (0, 1)} \left[ \left(g ^ {*}\right) ^ {\prime} (z) \right] w ^ {*} \cdot x \propto \frac {1}{d} w ^ {*} \cdot x. \tag {15}
$$
The learned feature is proportional to the true feature, so an appropriate choice of  $\overline{\eta}$  and choosing  $q = g^{*}$  in Theorem 1 implies that  $\| \mathbb{K}f^{*}\|_{L^2}^{-2} = d^2$  samples are needed to learn  $f^{*}$ .
4.2 Functions of quadratic features
The next example shows how three-layer networks can learn nonlinear features, and thus obtain a sample complexity improvement over two-layer networks.
Let  $\mathcal{X}_d = \mathcal{S}^{d-1}(\sqrt{d})$ , the sphere of radius  $\sqrt{d}$ , and  $\nu$  the uniform measure on  $\mathcal{X}_d$ . The integral operator  $\mathbb{K}$  has been well studied [39, 27, 38], and its eigenfunctions correspond to the spherical harmonics. Preliminaries on spherical harmonics and this eigendecomposition are given in Appendix F.
Consider the target  $f^{*}(x) = g^{*}(x^{T}Ax)$ , where  $A \in \mathbb{R}^{d \times d}$  is a symmetric matrix and  $g^{*}: \mathbb{R} \to \mathbb{R}$  is an unknown link function. In contrast to a single-index model, the feature  $x^{T}Ax$  we aim to learn is a quadratic function of  $x$ . Since one can write  $x^{T}Ax = x^{T}\big(A - \mathrm{Tr}(A)\cdot \frac{I}{d}\big)x + \mathrm{Tr}(A)$ , we without loss of generality assume  $\mathrm{Tr}(A) = 0$ . We also select the normalization  $\| A\| _F^2 = \frac{d + 2}{2d} = \Theta (1)$ ; this ensures that  $\mathbb{E}_{x\sim \nu}[(x^T Ax)^2 ] = 1$ . We first make the following assumptions on the target function.:
Assumption 5.  $\mathbb{E}_x[f^* (x)] = 0,\mathbb{E}_{z\sim \mathcal{N}(0,1)}[(g^*)'(z)] = \Theta (1),g^*$  is 1-Lipschitz, and  $(g^{*})''$  has polynomial growth.
The first assumption can be achieved via a preprocessing step which subtracts the mean of  $f^{*}$ , the second is a nondegeneracy condition, and the last two assume the target is sufficiently smooth.
We next require the eigenvalues of  $A$  to satisfy an incoherence condition:
Assumption 6. Define  $\kappa \coloneqq \| A\|_{op}\sqrt{d}$ . Then  $\kappa = o(\sqrt{d} /polylog(d))$
Note that  $\kappa \leq \sqrt{d}$ . If  $A$  has rank  $\Theta(d)$  and condition number  $\Theta(1)$ , then  $\kappa = \Theta(1)$ . Furthermore, when the entries of  $A$  are sampled i.i.d,  $\kappa = \tilde{\Theta}(1)$  with high probability by Wigner's semicircle law.
Finally, we make the following nondegeneracy assumption on the Gegenbauer decomposition of  $\sigma_{2}$  (defined in Appendix F). We show that  $\lambda_2^2 (\sigma_2) = O(d^{-2})$ , and later argue that the following assumption is mild and indeed satisfied by standard activations such as  $\sigma_{2} = \mathrm{ReLU}$ .
Assumption 7. Let  $\lambda_2(\sigma_2)$  be the 2nd Gegenbauer coefficient of  $\sigma_{2}$ . Then  $\lambda_2^2 (\sigma_2) = \Theta (d^{-2})$ .
We can verify Assumptions 1 to 4 hold for this setting, and thus applying Theorem 1 yields the following:
Theorem 3. Under Assumptions 5 to 7, with high probability Algorithm 1 satisfies the population loss bound
$$
\mathbb {E} _ {x} \left[ \left(f (x; \hat {\theta}) - f ^ {*} (x)\right) ^ {2} \right] \lesssim \tilde {O} \left(\frac {d ^ {4}}{\min  \left(n , m _ {1} , m _ {2}\right)} + \frac {1}{\sqrt {n}} + \left(\frac {\kappa}{\sqrt {d}}\right) ^ {1 / 3}\right) \tag {16}
$$
We thus require sample size  $n = \tilde{\Omega}(d^4)$  and widths  $m_1, m_2 = \tilde{\Omega}(d^4)$  to obtain  $o_d(1)$  test loss.
Proof Sketch. The integral operator  $\mathbb{K}$  has eigenspaces corresponding to spherical harmonics of degree  $k$ . In particular, [27] shows that, in  $L^2$ ,
$$
\mathbb {K} f ^ {*} = \sum_ {k \geq 0} c _ {k} P _ {k} f ^ {*}, \tag {17}
$$
where  $P_{k}$  is the orthogonal projection onto the subspace of degree  $k$  spherical harmonics, and the  $c_{k}$  are constants satisfying  $c_{k} = O(d^{-k})$ . Since  $f^{*}$  is an even function and  $\mathbb{E}_x[f^* (x)] = 0$ , truncating this expansion at  $k = 2$  yields
$$
\mathbb {K} f ^ {*} = \Theta (d ^ {- 2}) \cdot P _ {2} f ^ {*} + O (d ^ {- 4}), \tag {18}
$$
It thus suffices to compute  $P_{2}f^{*}$ . To do so, we draw a connection to the universality phenomenon in high dimensional probability. Consider two features  $x^{T}Ax$  and  $x^{T}Bx$  with  $\langle A,B\rangle = 0$ . We show that, when  $d$  is large, the distribution of  $x^{T}Ax$  approaches that of the standard Gaussian, while  $x^{T}Bx$  approaches a mixture of  $\chi^2$  and Gaussian random variables independent of  $x^{T}Ax$ . As such, we show
$$
\mathbb {E} _ {x} \left[ g ^ {*} \left(x ^ {T} A x\right) x ^ {T} B x \right] \approx \mathbb {E} _ {x} \left[ g ^ {*} \left(x ^ {T} A x\right) \right] \cdot \mathbb {E} _ {x} \left[ x ^ {T} B x \right] = 0 \tag {19}
$$
$$
\mathbb {E} _ {x} \left[ g ^ {*} \left(x ^ {T} A x\right) x ^ {T} A x \right] \approx \mathbb {E} _ {z \sim \mathcal {N} (0, 1)} \left[ g ^ {*} (z) z \right] = \mathbb {E} _ {z \sim \mathcal {N} (0, 1)} \left[ \left(g ^ {*}\right) ^ {\prime} (z) \right]. \tag {20}
$$
The second expression can be viewed as an approximate version of Stein's lemma, which was applied in Section 4.1 to compute the learned feature. Altogether, our key technical result (stated formally in Lemma 20) is that for  $f^{*} = g^{*}(x^{T}Ax)$ , the projection  $P_{2}f^{*}$  satisfies
$$
\left(P _ {2} f ^ {*}\right) (x) = \mathbb {E} _ {z \sim \mathcal {N} (0, 1)} \left[ \left(g ^ {*}\right) ^ {\prime} (z) \right] \cdot x ^ {T} A x + o _ {d} (1) \tag {21}
$$
The learned feature is thus  $\mathbb{K}f^{*}(x) = \Theta (d^{-2})\cdot (x^{T}Ax + o_{d}(1))$ ; plugging this into Theorem 1 yields the  $d^4$  sample complexity.
The full proof of Theorem 3 is deferred to Appendix E.2. In Appendix E.3 we show that when  $g^{*}$  is a degree  $p = O(1)$  polynomial, Algorithm 1 learns  $f^{*}$  in  $\tilde{O}(d^4)$  samples with an improved error floor.
Comparison to two-layer networks. Existing guarantees for two-layer networks cannot efficiently learn functions of the form  $f^{*}(x) = g^{*}(x^{T}Ax)$  for arbitrary Lipschitz  $g^{*}$ . In fact, in Section 5 we provide an explicit lower bound against two-layer networks efficiently learning a subclass of these functions. When  $g^{*}$  is a degree  $p$  polynomial, networks in the kernel regime require  $d^{2p} \gg d^4$  samples to learn  $f^{*}$  [27]. Improved guarantees for two-layer networks learn degree  $p'$  polynomials in  $r^{p'}$  samples when the target only depends on a rank  $r \ll d$  projection of the input [18]. However,  $g^{*}(x^{T}Ax)$  cannot be written in this form for some  $r \ll d$ , and thus existing guarantees do not apply. We conjecture that two-layer networks require  $d^{\Omega(p)}$  samples when  $g^{*}$  is a degree  $p$  polynomial.
Altogether, the ability of three-layer networks to efficiently learn the class of functions  $f^{*}(x) = g^{*}(x^{T}Ax)$  hinges on their ability to extract the correct nonlinear feature. Empirical validation of the above examples is given in Appendix A.
5 An Optimization-Based Depth Separation
We complement the learning guarantee in Section 4.2 with a lower bound showing that there exist functions in this class that cannot be approximated by a polynomial size two-layer network.
The class of candidate two-layer networks is as follows. For a parameter vector  $\theta = (a, W, b_1, b_2)$ , where  $a \in \mathbb{R}^m$ ,  $W \in \mathbb{R}^{m \times d}$ ,  $b_1 \in \mathbb{R}^m$ ,  $b_2 \in \mathbb{R}$ , define the associated two-layer network as
$$
N _ {\theta} (x) := a ^ {T} \sigma \left(W x + b _ {1}\right) + b _ {2} = \sum_ {i = 1} ^ {m} a _ {i} \sigma \left(w _ {i} ^ {T} x + b _ {1, i}\right) + b _ {2}. \tag {22}
$$
Let  $\| \theta \|_{\infty} \coloneqq \max(\| a \|_{\infty}, \| W \|_{\infty}, \| b_1 \|_{\infty}, \| b_2 \|_{\infty})$  denote the maximum parameter value. We make the following assumption on  $\sigma$ , which holds for all commonly-used activations.
Assumption 8. There exist constants  $C_{\sigma}, \alpha_{\sigma}$  such that  $|\sigma(z)| \leq C_{\sigma}(1 + |z|)^{\alpha_{\sigma}}$ .
Our main theorem establishing the separation is the following.
Theorem 4. Let  $d$  be a sufficiently large even integer. Consider the target function  $f^{*}(x) = \mathrm{ReLU}(x^{T}Ax) - c_{0}$ , where  $A = \frac{1}{\sqrt{d}} U\left( \begin{array}{cc}0 & I_{d / 2}\\ I_{d / 2} & 0 \end{array} \right)U^{T}$  for some orthogonal matrix  $U$  and  $c_{0} = \mathbb{E}_{x\sim \nu}\left[\mathrm{ReLU}(x^T Ax)\right]$ . Under Assumption 8, there exist constants  $C_1,C_2,C_3,c_3$  depending only on  $(C_{\sigma},\alpha_{\sigma})$ , such that for any  $c_{3}\geq \epsilon \geq C_{3}d^{-2}$ , any two layer neural network  $N_{\theta}(x)$  of width  $m$  and population  $L^2$  error bound  $\| N_{\theta} - f^{*}\|_{L^2}^2\leq \epsilon$  must satisfy  $\max (m,\| \theta \|_{\infty})\geq C_1\exp \big(C_2\epsilon^{-1 / 2}\log (d\epsilon)\big)$ . However, Algorithm 1 outputs a predictor satisfying the population  $L^2$  loss bound
$$
\mathbb {E} _ {x} \left[ \left(f (x; \hat {\theta}) - f ^ {*} (x)\right) ^ {2} \right] \lesssim O \left(\frac {d ^ {4}}{n} + \sqrt {\frac {d}{n}} + d ^ {- 1 / 6}\right). \tag {23}
$$
after  $T = \mathrm{poly}(d, m_1, m_2, n)$  timesteps.
The lower bound follows from a modification of the argument in [19] along with an explicit decomposition of the ReLU function into spherical harmonics. We remark that the separation applies for any link function  $g^{*}$  whose Gegenbauer coefficients decay sufficiently slow. The upper bound follows from an application of Theorem 1 to a smoothed version of ReLU, since ReLU is not in  $\mathcal{W}^{2,\infty}([-1,1])$ . The full proof of the theorem is given in Appendix G.
Remarks. In order for a two-layer network to achieve test loss matching the  $d^{-1/6}$  error floor, either the width or maximum weight norm of the network must be  $\exp\bigl(\Omega(d^{\delta})\bigr)$  for some constant  $\delta$ ; this lower bound is superpolynomial in  $d$ . As a consequence, gradient descent on a poly-width two-layer neural network with stable step size must run for superpolynomially many iterations in order for some weight to grow this large and thus converge to a solution with low test error. Therefore  $f^{*}$  is not learnable via gradient descent in polynomial time. This reduction from a weight norm lower bound to a runtime lower bound is made precise in [45].
We next describe a specific example of such an  $f^{*}$ . Let  $S$  be a  $d/2$ -dimensional subspace of  $\mathbb{R}^d$ , and let  $A = d^{-\frac{1}{2}}P_S - d^{-\frac{1}{2}}P_S^\perp$ , where  $P_S, P_S^\perp$  are projections onto the subspace  $S$  and its orthogonal complement respectively. Then,  $f^{*}(x) = \mathrm{ReLU}\left(2d^{-\frac{1}{2}}\| P_S x\|^2 - d^{\frac{1}{2}}\right)$ . [44] established an optimization-based separation for the target  $\mathrm{ReLU}(1 - \| x \|)$ , under a different distribution  $\nu$ . However, their analysis relies heavily on the rotational symmetry of the target, and they posed the question of learning  $\mathrm{ReLU}(1 - \| P_S x \|)$  for some subspace  $S$ . Our separation applies to a similar target function, and crucially does not rely on this rotational invariance.
In this work we showed that three-layer networks can both learn nonlinear features and leverage these features to obtain a provable sample complexity improvement over two-layer networks. There are a number of interesting directions for future work. First, can our framework be used to learn hierarchical functions of a larger class of features beyond quadratic functions? Next, since Theorem 1 is general
purpose and makes minimal distributional assumptions, it would be interesting to understand if it can be applied to standard empirical datasets such as CIFAR-10, and what the learned feature  $\mathbb{K}f^{*}$  and hierarchical learning correspond to in this setting. Finally, our analysis studies the nonlinear feature learning that arises from a neural representation  $\sigma_{2}(Vx)$  where  $V$  is fixed at initialization. This alone was enough to establish a separation between two and three-layer networks. A fascinating question is to understand the additional features that can be learned when both  $V$  and  $W$  are jointly trained. Such an analysis, however, is incredibly challenging in feature-learning regimes.
EN acknowledges support from a National Defense Science & Engineering Graduate Fellowship. AD acknowledges support from a NSF Graduate Research Fellowship. EN, AD, and JDL acknowledge support of the ARO under MURI Award W911NF-11-1-0304, the Sloan Research Fellowship, NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award, and NSF CAREER Award 2144994. The authors would like to thank Itay Safran for helpful discussions.
[1] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Proceedings of Thirty Fifth Conference on Learning Theory, pages 4782-4887, 2022.
[2] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In *The Thirty Sixth Annual Conference on Learning Theory*, pages 2552–2623. PMLR, 2023.
[3] Zeyuan Allen-Zhu and Yanzhi Li. What can resnet learn efficiently, going beyond kernels? Advances in Neural Information Processing Systems, 32, 2019.
[4] Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04413, 2020.
[5] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in neural information processing systems, 32, 2019.
[6] Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[7] Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning (ICML), 2019.
[8] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
[9] Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. In International Conference on Learning Representations, 2020.
[10] Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[11] William Beckner. Sobolev inequalities, the poisson semigroup, and analysis on the sphere sn. Proceedings of the National Academy of Sciences of the United States of America, 89 11: 4816-9, 1992.
One limitation of our work is that the innermost layer weights  $V$  are fixed throughout training, and our training procedure is layerwise. We remark that this is already enough to establish an algorithmic separation between two and three-layer networks (Section 5), and that prior works on two-layer networks also rely on layerwise training procedures. It is a very interesting direction of future work to understand if training  $V$  can lead to learning a larger class of functions. Another limitation is that our examples in Section 4 rely on the data distribution  $\nu$  being either the standard Gaussian or uniform on the sphere. This allows us to carefully characterize the integral operator  $\mathbb{K}$ , and we remark that other works on two-layer networks make similar distributional assumptions (such as uniform on the hypercube). However, it is a very important direction of future work to extend these lines of work to more general data distributions.
A Empirical Validation
Single Index Setting
Quadratic Feature (Random Projection)
Figure 1: We ran Algorithm 1 on both the single index and quadratic feature settings described in Section 4. Each trial was run with 5 random seeds. The solid lines represent the medians and the shaded areas represent the min and max values. For every trial we recorded both the test loss on a test set of size  $2^{15}$  and the linear correlation between the learned feature map  $\phi(x)$  and the true intermediate feature  $h^{\star}(x)$  where  $h^{\star}(x) = x \cdot \beta$  for the single index setting and  $h^{\star}(x) = x^{T}Ax$  for the quadratic feature setting. Our results show that the test loss goes to 0 as the linear correlation between the learned feature map  $\phi$  and the true intermediate feature  $h^{\star}$  approaches 1.
Quadratic Feature (Random Entries)
Correlation Between Learned and Predicted Features
We empirically verify our conclusions in the single index setting of Section 4.1 and the quadratic feature setting of Section 4.2:
Single Index Setting We learn the target function  $g^{\star}(w^{\star} \cdot x)$  using Algorithm 1 where  $w^{\star} \in S^{d-1}$  is drawn randomly and  $g^{\star}(x) = \operatorname{sigmoid}(x) = \frac{1}{1 + e^{-x}}$ , which satisfies the condition  $\mathbb{E}_{z \sim \mathcal{N}(0,1)}[g'(z)] \neq 0$ . As in Theorem 2, we choose the initial activation  $\sigma_2(z) = z$ . We optimize the hyperparameters  $\eta_1, \lambda$  using grid search over a holdout validation set of size  $2^{15}$  and report the final error over a test set of size  $2^{15}$ .
Quadratic Feature Setting We learn the target function  $g^{\star}(x^{T}Ax)$  using Algorithm 1 where  $g^{\star}(x) = \mathrm{ReLU}(x)$ . We ran our experiments with two different choices of  $A$ :
-  $A$  is symmetric with random entries, i.e.  $A_{ij} \sim N(0,1)$  and  $A_{ji} = A_{ij}$  for  $i \leq j$ .
-  $A$  is a random projection, i.e.  $A = \Pi_S$  where  $S$  is a random  $d/2$  dimensional subspace.
Both choices of  $A$  were then normalized so that  $\operatorname{tr} A = 0$  and  $\| A \|_F = 1$  by subtracting the trace and dividing by the Frobenius norm. We chose initial activation  $\sigma_2(z) = \mathrm{ReLU}(z)$ . We note that in both examples,  $\kappa = \Theta(1)$ . As above, we optimize the hyperparameters  $\eta_1, \lambda$  using grid search over a holdout validation set of size  $2^{15}$  and report the final error over a test set of size  $2^{15}$ .