SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
L (w) = \sum_ {k \geq 0} \frac {c _ {k} ^ {2}}{k !} [ 1 - \alpha^ {k} ] \quad \text {a n d} \quad \nabla L (w) = - \left(P _ {w} ^ {\perp} w ^ {\star}\right) \sum_ {k \geq 0} \frac {c _ {k} ^ {2}}{(k - 1) !} \alpha^ {k - 1}.
$$
As  $\alpha \ll 1$  throughout most of the trajectory, the gradient is dominated by the first nonzero Hermite coefficient so up to constants:  $\mathbb{E}[v\cdot w^{\star}] = -\nabla L(w)\cdot w^{\star}\approx \alpha^{k^{\star} - 1}$ . Similarly, a standard concentration argument shows that because  $v$  is a random vector in  $d$  dimensions where each coordinate is  $O(1)$ ,  $\mathbb{E}[\| v\| ]^2\approx d$ . Therefore the SNR is equal to  $\alpha^{2(k^{\star} - 1)} / d$  so with an optimal learning rate schedule,
$$
\mathbb {E} \left[ \alpha_ {t + 1} \mid \mathcal {F} _ {t} \right] \approx \alpha_ {t} + \alpha_ {t} ^ {2 k ^ {*} - 3} / d.
$$
This can be approximated by the ODE  $\alpha' = \alpha^{2k^{\star} - 3} / d$ . Solving this ODE with the initial  $\alpha_0 \asymp d^{-1/2}$  gives that the escape time is proportional to  $d\alpha_0^{-2(k^{\star} - 1)} = d^{k^{\star} - 1}$  which heuristically re-derives the result of Ben Arous et al. [1].
5.3 How Smoothing boosts the SNR
Smoothing improves the sample complexity of online SGD by boosting the SNR of the stochastic gradient  $\nabla L(w;x;y)$ .
Figure 1: When  $\lambda = d^{1/4}$ , smoothing increases the SNR until  $\alpha = \lambda d^{-1/2} = d^{-1/4}$ .
Computing the Signal Recall that the population loss was approximately equal to  $1 - \frac{c_{k^{\star}}^{2}}{k^{\star!}}\alpha^{k^{\star}}$  where  $k^{\star}$  is the first nonzero Hermite coefficient of  $\sigma$ . Isolating the dominant  $\alpha^{k^{\star}}$  term and applying the smoothing operator  $\mathcal{L}_{\lambda}$ , we get:
$$
\mathcal {L} _ {\lambda} \left(\alpha^ {k ^ {\star}}\right) = \mathbb {E} _ {z \sim \mu_ {w}} \left[ \left(\frac {w + \lambda z}{\| w + \lambda z \|} \cdot w ^ {\star}\right) ^ {k ^ {\star}} \right].
$$
Because  $z \perp w$  and  $\| z \| = 1$  we have that  $\| w + \lambda z \| = \sqrt{1 + \lambda^2} \approx \lambda$ . Therefore,
$$
\mathcal {L} _ {\lambda} (\alpha^ {k ^ {\star}}) \approx \lambda^ {- k ^ {\star}}   \mathbb {E} _ {z \sim \mu_ {w}} \left[ (\alpha + \lambda (z \cdot w ^ {\star})) ^ {k ^ {\star}} \right] = \lambda^ {- k ^ {\star}} \sum_ {j = 0} ^ {k} \binom {k ^ {\star}} {j} \alpha^ {k ^ {\star} - j} \lambda^ {j}   \mathbb {E} _ {z \sim \mu_ {w}} [ (z \cdot w ^ {\star}) ^ {j} ].
$$
Now because  $z \stackrel{d}{=} -z$ , the terms where  $j$  is odd disappear. Furthermore, for a random  $z$ ,  $|z \cdot w^{\star}| = \Theta(d^{-1/2})$ . Therefore reindexing and ignoring all constants we have that
$$
L _ {\lambda} (w) \approx 1 - \mathcal {L} _ {\lambda} (\alpha^ {k ^ {\star}}) \approx 1 - \lambda^ {- k ^ {\star}} \sum_ {j = 0} ^ {\lfloor \frac {k ^ {\star}}{2} \rfloor} \alpha^ {k ^ {\star} - 2 j} \left(\lambda^ {2} / d\right) ^ {j}.
$$
Differentiating gives that
$$
\mathbb {E} [ v \cdot w ^ {\star} ] \approx - w ^ {\star} \cdot \nabla_ {w} \mathcal {L} _ {\lambda} (\alpha^ {k ^ {\star}}) \approx \lambda^ {- 1} \sum_ {j = 0} ^ {\lfloor \frac {k ^ {\star} - 1}{2} \rfloor} \left(\frac {\alpha}{\lambda}\right) ^ {k ^ {\star} - 1} \left(\frac {\lambda^ {2}}{\alpha^ {2} d}\right) ^ {j}.
$$
As this is a geometric series, it is either dominated by the first or the last term depending on whether  $\alpha \geq \lambda d^{-1/2}$  or  $\alpha \leq \lambda d^{-1/2}$ . Furthermore, the last term is either  $d^{-\frac{k^{\star} - 1}{2}}$  if  $k^{\star}$  is odd or  $\frac{\alpha}{\lambda} d^{-\frac{k^{\star} - 2}{2}}$  if  $k^{\star}$  is even. Therefore the signal term is:
$$
\mathbb {E} [ v \cdot w ^ {\star} ] \approx \lambda^ {- 1} \left\{ \begin{array}{l l} (\frac {\alpha}{\lambda}) ^ {k ^ {\star} - 1} & \alpha \geq \lambda d ^ {- 1 / 2} \\ d ^ {- \frac {k ^ {\star} - 1}{2}} & \alpha \leq \lambda d ^ {- 1 / 2} \text {a n d} k ^ {\star} \text {i s o d d} \\ \frac {\alpha}{\lambda} d ^ {- \frac {k ^ {\star} - 2}{2}} & \alpha \leq \lambda d ^ {- 1 / 2} \text {a n d} k ^ {\star} \text {i s e v e n} \end{array} . \right.
$$
Computing the Noise Recall that  $L_{\lambda}(w; x, y) = 1 - y \sigma(w \cdot x)$ . Differentiating through the smoothing operator gives:
$$
\nabla_ {w} L _ {\lambda} (w; x, y) = - y \nabla_ {w} \mathcal {L} _ {\lambda} (\sigma (w \cdot x)) \approx \lambda^ {- 1} y x \mathcal {L} _ {\lambda} \left(\sigma^ {\prime} (w \cdot x)\right).
$$
We know that with high probability,  $y = \tilde{O}(1)$  and  $\| x \| = O(\sqrt{d})$  so it suffices to bound  $\mathcal{L}_{\lambda}(\sigma'(w \cdot x))$ . The variance of this term is equal to:
$$
\mathbb {E} _ {x} [ \mathcal {L} _ {\lambda} (\sigma^ {\prime} (w \cdot x)) ^ {2} ] = \mathbb {E} _ {x} \left[ \mathbb {E} _ {z \sim \mu_ {w}} \left[ \sigma^ {\prime} \left(\frac {w + \lambda z}{\sqrt {1 + \lambda^ {2}}} \cdot x\right) \right] ^ {2} \right].
$$
To compute this expectation, we will create an i.i.d. copy  $z'$  of  $z$  and rewrite this expectation as:
$$
\mathbb {E} _ {x} [ \mathcal {L} _ {\lambda} (\sigma^ {\prime} (w \cdot x)) ^ {2} ] = \mathbb {E} _ {x} \left[ \mathbb {E} _ {z, z ^ {\prime} \sim \mu_ {w}} \left[ \sigma^ {\prime} \left(\frac {w + \lambda z}{\sqrt {1 + \lambda^ {2}}} \cdot x\right) \sigma^ {\prime} \left(\frac {w + \lambda z ^ {\prime}}{\sqrt {1 + \lambda^ {2}}} \cdot x\right) \right] \right].
$$
Now we can swap the expectations and compute the expectation with respect to  $x$  first using the Hermite expansion of  $\sigma$ . As the first nonzero Hermite coefficient of  $\sigma'$  is  $k^{\star} - 1$ , this variance is approximately equal to the correlation between  $\frac{w + \lambda z}{\sqrt{1 + \lambda^2}}$  and  $\frac{w + \lambda z'}{\sqrt{1 + \lambda^2}}$  raised to the  $k^{\star} - 1$  power:
$$
\mathbb {E} _ {x} [ \mathcal {L} _ {\lambda} (\sigma^ {\prime} (w \cdot x)) ^ {2} ] \approx \mathbb {E} _ {z, z ^ {\prime} \sim \mu_ {w}} \left[ \left(\frac {w + \lambda z}{\sqrt {1 + \lambda^ {2}}} \cdot \frac {w + \lambda z ^ {\prime}}{\sqrt {1 + \lambda^ {2}}}\right) ^ {k ^ {\star} - 1} \right] = \mathbb {E} _ {z, z ^ {\prime} \sim \mu_ {w}} \left[ \left(\frac {1 + \lambda^ {2} z \cdot z ^ {\prime}}{1 + \lambda^ {2}}\right) ^ {k ^ {\star} - 1} \right]
$$
As  $z, z'$  are random unit vectors, their inner product is of order  $d^{-1/2}$ . Therefore when  $\lambda \ll d^{1/4}$ , the first term in the numerator is dominant and when  $\lambda \gg d^{1/4}$ , the second term is dominant. Combining these two regimes gives that the variance is of order  $\min(\lambda, d^{1/4})^{k^{\star} - 1}$  which motivates our choice of  $\lambda = d^{1/4}$ . Combining this with  $y = \tilde{O}(1)$  and  $\|x\| = O(\sqrt{d})$  gives that for  $\lambda \leq d^{1/4}$ , the variance of the noise is bounded by  $\mathbb{E}[\|v\|^2] \lesssim d\lambda^{-2k^{\star}}$ . Note that in the high signal regime ( $\alpha \geq \lambda d^{-1/2}$ ), both the signal and the noise are smaller by factors of  $\lambda^{k^{\star}}$  which cancel when computing the SNR. However, when  $\alpha \leq \lambda d^{-1/2}$  the smoothing shrinks the noise faster than it shrinks the signal, resulting in an overall larger SNR. Explicitly,
$$
\mathrm {S N R} := \frac {\mathbb {E} [ v \cdot w ^ {\star} ] ^ {2}}{\mathbb {E} [ \| v \| ^ {2} ]} \approx \frac {1}{d} \left\{ \begin{array}{l l} \alpha^ {2 (k ^ {\star} - 1)} & \alpha \geq \lambda d ^ {- 1 / 2} \\ (\lambda^ {2} / d) ^ {k ^ {\star} - 1} & \alpha \leq \lambda d ^ {- 1 / 2} \text {a n d} k ^ {\star} \text {i s o d d} \\ \alpha^ {2} (\lambda^ {2} / d) ^ {k ^ {\star} - 2} & \alpha \leq \lambda d ^ {- 1 / 2} \text {a n d} k ^ {\star} \text {i s e v e n} \end{array} . \right.
$$
For  $\alpha \geq \lambda d^{-1/2}$ , smoothing does not affect the SNR. However, when  $\alpha \leq \lambda d^{-1/2}$ , smoothing greatly increases the SNR (see Figure 1).
Solving the ODE:  $\alpha' = \mathrm{SNR} / \alpha$  gives that it takes  $T \approx d^{k^{\star} - 1}\lambda^{-2k^{\star} + 4}$  steps to converge to  $\alpha \approx 1$  from  $\alpha \approx d^{-1/2}$ . Once  $\alpha \approx 1$ , the problem is locally strongly convex, so we can decay the learning rate and use classical analysis of strongly-convex functions to show that  $\alpha \geq 1 - \epsilon$  with an additional  $d/\epsilon$  steps, from which Theorem 1 follows.
For  $k^{\star} = 3,4,5$  and  $d = 2^{8},\ldots ,2^{13}$  we ran a minibatch variant of Algorithm 1 with batch size  $B$  when  $\sigma (x) = \frac{He_{k^{\star}}(x)}{\sqrt{k!}}$ , the normalized  $k^{\star}$ th Hermite polynomial. We set:
$$
\lambda = d ^ {1 / 4}, \quad \eta = \frac {B d ^ {- k ^ {\star} / 2} (1 + \lambda^ {2}) ^ {k ^ {\star} - 1}}{1 0 0 0 k ^ {\star} !}, \quad B = \min  \left(0. 1 d ^ {k ^ {\star} / 2} (1 + \lambda^ {2}) ^ {- 2 k ^ {\star} + 4}, 8 1 9 2\right).
$$
We computed the number of samples required for Algorithm 1 to reach  $\alpha^2 = 0.5$  from  $\alpha = d^{-1/2}$  and we report the min, mean, and max over 10 random seeds. For each  $k$  we fit a power law of the form  $n = c_1 d^{c_2}$  in order to measure how the sample complexity scales with  $d$ . For all values of  $k^\star$ , we find that  $c_2 \approx k^\star / 2$  which matches Theorem 1. The results can be found in Figure 2 and additional experimental details can be found in Appendix E.
7.1 Tensor PCA
We next outline connections to the Tensor PCA problem. Introduced in [8], the goal of Tensor PCA is to recover the hidden direction  $w^{*} \in S^{d - 1}$  from the noisy  $k$ -tensor  $T_{n} \in (\mathbb{R}^{d})^{\otimes k}$  given by
$$
T _ {n} = (w ^ {*}) ^ {\otimes k} + \frac {1}{\sqrt {n}} Z,
$$
where  $Z \in (\mathbb{R}^d)^{\otimes k}$  is a Gaussian noise tensor with each entry drawn i.i.d from  $\mathcal{N}(0,1)$ .
The Tensor PCA problem has garnered significant interest as it exhibits a statistical-computational gap.  $w^{*}$  is information theoretically recoverable when  $n \gtrsim d$ . However, the best polynomial-time algorithms require  $n \gtrsim d^{k/2}$ ; this lower bound has been shown to be tight for various notions of
This normalization is equivalent to the original  $\frac{1}{\beta\sqrt{d}} Z$  normalization by setting  $n = \beta^2 d$ .
Figure 2: For  $k^{\star} = 3,4,5$ , Algorithm 1 finds  $w^{\star}$  with  $n \propto d^{k^{\star} / 2}$  samples. The solid lines and the shaded areas represent the mean and min/max values over 10 random seeds. For each curve, we also fit a power law  $n = c_{1}d^{c_{2}}$  represented by the dashed lines. The value of  $c_{2}$  is reported in the legend.
hardness such as CSQ or SoS lower bounds [34, 25, 35-39]. Tensor PCA also exhibits a gap between spectral methods and iterative algorithms. Algorithms that work in the  $n \asymp d^{k/2}$  regime rely on unfolding or contracting the tensor  $X$ , or on semidefinite programming relaxations [34, 25]. On the other hand, iterative algorithms including gradient descent, power method, and AMP require a much larger sample complexity of  $n \gtrsim d^{k-1}$  [40]. The suboptimality of iterative algorithms is believed to be due to bad properties of the landscape of the Tensor PCA objective in the region around the initialization. Specifically [41, 42] argue that there are exponentially many local minima near the equator in the  $n \ll d^{k-1}$  regime. To overcome this, prior works have considered "smoothed" versions of gradient descent, and show that smoothing recovers the computationally optimal SNR in the  $k = 3$  case [26] and heuristically for larger  $k$  [27].
7.1.1 The Partial Trace Algorithm
The smoothing algorithms above are inspired by the following partial trace algorithm for Tensor PCA [25], which can be viewed as Algorithm 1 in the limit as  $\lambda \to \infty$  [26]. Let  $T_{n} = (w^{\star})^{\otimes k} + \frac{1}{\sqrt{n}} Z$ . Then we will consider iteratively contracting indices of  $T$  until all that remains is a vector (if  $k$  is odd) or a matrix (if  $k$  is even). Explicitly, we define the partial trace tensor by
$$
M _ {n} := T _ {n} \Big (I _ {d} ^ {\otimes \lceil \frac {k - 2}{2} \rceil} \Big) \in \left\{ \begin{array}{l l} \mathbb {R} ^ {d \times d} & k \text {i s e v e n} \\ \mathbb {R} ^ {d} & k \text {i s o d d}. \end{array} \right.
$$
When  $k^{\star}$  is odd, we can directly return  $M_{n}$  as our estimate for  $w^{\star}$  and when  $k^{\star}$  is even we return the top eigenvector of  $M_{n}$ . A standard concentration argument shows that this succeeds when  $n \gtrsim d^{[k/2]}$ . Furthermore, this can be strengthened to  $d^{k/2}$  by using the partial trace vector as a warm start for gradient descent or tensor power method when  $k$  is odd [26, 27].
7.1.2 The Connection Between Single Index Models and Tensor PCA
For both tensor PCA and learning single index models, gradient descent succeeds when the sample complexity is  $n = d^{k - 1}$  [40, 1]. On the other hand, the smoothing algorithms for Tensor PCA [27, 26] succeed with the computationally optimal sample complexity of  $n = d^{k / 2}$ . Our Theorem 1 shows that this smoothing analysis can indeed be transferred to the single-index model setting.
In fact, one can make a direct connection between learning single-index models with Gaussian covariates and Tensor PCA. Consider learning a single-index model when  $\sigma(x) = \frac{He_k(x)}{\sqrt{k!}}$ , the normalized  $k$ th Hermite polynomial. Then minimizing the correlation loss is equivalent to maximizing
the loss function:
$$
L _ {n} (w) = \frac {1}{n} \sum_ {i = 1} ^ {n} y _ {i} \frac {H e _ {k} (w \cdot x _ {i})}{\sqrt {k !}} = \left\langle w ^ {\otimes k}, T _ {n} \right\rangle \quad \text {w h e r e} \quad T _ {n} := \frac {1}{n} \sum_ {i = 1} ^ {n} y _ {i} \frac {\mathbf {H e} _ {k} (x _ {i})}{\sqrt {k !}}.
$$
Here  $\mathbf{H}\mathbf{e}_k(x_i)\in (\mathbb{R}^d)^{\otimes k}$  denotes the  $k$ th Hermite tensor (see Appendix A.2 for background on Hermite polynomials and Hermite tensors). In addition, by the orthogonality of the Hermite tensors,  $\mathbb{E}_{x,y}[T_n] = (w^\star)^{\otimes k}$  so we can decompose  $T_{n} = (w^{\star})^{\otimes k} + Z_{n}$  where by standard concentration, each entry of  $Z_{n}$  is order  $n^{-1 / 2}$ . We can therefore directly apply algorithms for Tensor PCA to this problem. We remark that this connection is a heuristic, as the structure of the noise in Tensor PCA and our single index model setting are different.
7.2 Empirical Risk Minimization on the Smoothed Landscape
Our main sample complexity guarantee, Theorem 1, is based on a tight analysis of online SGD (Algorithm 1) in which each sample is used exactly once. One might expect that if the algorithm were allowed to reuse samples, as is standard practice in deep learning, that the algorithm could succeed with fewer samples. In particular, Abbe et al. [4] conjectured that gradient descent on the empirical loss  $L_{n}(w)\coloneqq \frac{1}{n}\sum_{i = 1}^{n}L(w;x_{i};y_{i})$  would succeed with  $n\gtrsim d^{k^{\star} / 2}$  samples.
Our smoothing algorithm Algorithm 1 can be directly translated to the ERM setting to learn  $w^{\star}$  with  $n \gtrsim d^{k^{\star}/2}$  samples. We can then Taylor expand the smoothed loss in the large  $\lambda$  limit:
$$
\mathcal {L} _ {\lambda} (L _ {n} (w)) \approx \mathbb {E} _ {z \sim \mathrm {U n i f} (S ^ {d - 1})} \left[ L _ {n} (z) + \lambda^ {- 1} w \cdot \nabla L _ {n} (z) + \frac {\lambda^ {- 2}}{2} \cdot w ^ {T} \nabla^ {2} L _ {n} (z) w \right] + O (\lambda^ {- 3}).
$$
As  $\lambda \to \infty$ , gradient descent on this smoothed loss will converge to  $\mathbb{E}_{z\sim \mathrm{Unif}(S^{d - 1})}[\nabla L_n(z)]$  which is equivalent to the partial trace estimator for  $k^{\star}$  odd (see Section 7.1). If  $k^{\star}$  is even, this first term is zero in expectation and gradient descent will converge to the top eigenvector of  $\mathbb{E}_{z\sim \mathrm{Unif}(S^{d - 1})}[\nabla^2 L_n(z)]$ , which corresponds to the partial trace estimator for  $k^{\star}$  even. Mirroring the calculation for the partial trace estimator, this succeeds with  $n\gtrsim d^{[k^{\star} / 2]}$  samples. When  $k^{\star}$  is odd, this can be further improved to  $d^{k^{\star} / 2}$  by using this estimator as a warm start from which to run gradient descent with  $\lambda = 0$  as in Anandkumar et al. [26], Biroli et al. [27].
7.3 Connection to Minibatch SGD
A recent line of works has studied the implicit regularization effect of stochastic gradient descent [9, 11, 10]. The key idea is that over short timescales, the iterates converge to a quasi-stationary distribution  $N(\theta, \lambda S)$  where  $S \approx I$  depends on the Hessian and the noise covariance at  $\theta$  and  $\lambda$  is proportional to the ratio of the learning rate and batch size. As a result, over longer periods of time SGD follows the smoothed gradient of the empirical loss:
$$
\widetilde {L} _ {n} (w) = \mathbb {E} _ {z \sim N (0, S)} [ L _ {n} (w + \lambda z) ].
$$
We therefore conjecture that minibatch SGD is also able to achieve the optimal  $n \gtrsim d^{k^{\star} / 2}$  sample complexity without explicit smoothing if the learning rate and batch size are properly tuned.
AD acknowledges support from a NSF Graduate Research Fellowship. EN acknowledges support from a National Defense Science & Engineering Graduate Fellowship. RG is supported by NSF Award DMS-2031849, CCF-1845171 (CAREER), CCF-1934964 (Tripods) and a Sloan Research Fellowship. AD, EN, 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.
[1] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research, 22(1):4788–4838, 2021.
[35] Dmitriy Kunisky, Alexander S. Wein, and Afonso S. Bandeira. Notes on computational hardness of hypothesis testing: Predictions using the low-degree likelihood ratio, 2019.
[36] Afonso S Bandeira, Ahmed El Alaoui, Samuel Hopkins, Tselil Schramm, Alexander S Wein, and Ilias Zadik. The Franz-parisi criterion and computational trade-offs in high dimensional statistics. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 33831-33844. Curran Associates, Inc., 2022.
[37] Matthew S Brennan, Guy Bresler, Sam Hopkins, Jerry Li, and Tselil Schramm. Statistical query algorithms and low degree tests are almost equivalent. In Mikhail Belkin and Samory Kpotufe, editors, Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages 774-774. PMLR, 15-19 Aug 2021. URL https://proceedings.mlr.press/v134/brennan21a.html.
[38] Rishabh Dudeja and Daniel Hsu. Statistical query lower bounds for tensor pca. Journal of Machine Learning Research, 22(83):1-51, 2021. URL http://jmlr.org/papers/v22/20-837.html.
[39] Rishabh Dudeja and Daniel Hsu. Statistical-computational trade-offs in tensor pca and related problems via communication complexity, 2022.
[40] Gérard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Algorithmic thresholds for tensor PCA. The Annals of Probability, 48(4):2052 - 2087, 2020. doi: 10.1214/19-AOP1415. URL https://doi.org/10.1214/19-AOP1415.
[41] Valentina Ros, Gerard Ben Arous, Giulio Biroli, and Chiara Cammarota. Complex energy landscapes in spiked-tensor and simple glassy models: Ruggedness, arrangements of local minima, and phase transitions. Phys. Rev. X, 9:011003, Jan 2019. doi: 10.1103/PhysRevX.9.011003. URL https://link.aps.org/doi/10.1103/PhysRevX.9.011003.
[42] Gérard Ben Arous, Song Mei, Andrea Montanari, and Mihai Nica. The landscape of the spiked tensor model. Communications on Pure and Applied Mathematics, 72(11):2282-2330, 2019. doi: https://doi.org/10.1002/cpa.21861. URL https://onlinelibrary.wiley.com/doi/abs/10.1002/cpa.21861.
[43] Balázs Szö rényi. Characterizing statistical query learning: Simplified notions and proofs. In ALT, 2009.
[44] Iosif Pinelis. Optimum Bounds for the Distributions of Martingales in Banach Spaces. The Annals of Probability, 22(4):1679 - 1706, 1994. doi: 10.1214/aop/1176988477. URL https://doi.org/10.1214/aop/1176988477.
[45] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax.
[46] Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https://www.wandb.com/. Software available from wandb.com.
A Background and Notation
A.1 Tensor Notation
Throughout this section let  $T \in (\mathbb{R}^d)^{\otimes k}$  be a  $k$ -tensor.
Definition 5 (Tensor Action). For a  $j$  tensor  $A \in (\mathbb{R}^d)^{\otimes j}$  with  $j \leq k$ , we define the action  $T[A]$  of  $T$  on  $A$  by
$$
(T [ A ]) _ {i _ {1}, \dots , i _ {k - j}} := T _ {i _ {1}, \dots , i _ {k}} A ^ {i _ {k - j + 1}, \dots , i _ {k}} \in (\mathbb {R} ^ {d}) ^ {\otimes (k - j)}.
$$
We will also use  $\langle T, A \rangle$  to denote  $T[A] = A[T]$  when  $A, T$  are both  $k$  tensors. Note that this corresponds to the standard dot product after flattening  $A, T$ .
Definition 6 (Permutation/Transposition). Given a  $k$ -tensor  $T$  and a permutation  $\pi \in S_k$ , we use  $\pi(T)$  to denote the result of permuting the axes of  $T$  by the permutation  $\pi$ , i.e.
$$
\pi (T) _ {i _ {1}, \dots , i _ {k}} := T _ {i _ {\pi (1)}, \dots , i _ {\pi (k)}}.
$$
Definition 7 (Symmetrization). We define  $\mathrm{Sym}_k\in (\mathbb{R}^d)^{\otimes 2k}$  by
$$
\left(\operatorname {S y m} _ {k}\right) _ {i _ {1}, \dots , i _ {k}, j _ {1}, \dots , j _ {k}} = \frac {1}{k !} \sum_ {\pi \in S _ {k}} \delta_ {i _ {\pi (1)}, j _ {1}} \dots \delta_ {i _ {\pi (k)}, j _ {k}}
$$
where  $S_{k}$  is the symmetric group on  $1,\ldots ,k$ . Note that  $\mathrm{Sym}_k$  acts on  $k$  tensors  $T$  by
$$
(\operatorname {S y m} _ {k} [ T ]) _ {i _ {1}, \dots , i _ {k}} = \frac {1}{k !} \sum_ {\pi \in S _ {k}} \pi (T).
$$
i.e.  $\operatorname{Sym}_k[T]$  is the symmetrized version of  $T$ .
We will also overload notation and use  $\operatorname{Sym}$  to denote the symmetrization operator, i.e. if  $T$  is a  $k$ -tensor,  $\operatorname{Sym}(T) \coloneqq \operatorname{Sym}_k[T]$ .
Definition 8 (Symmetric Tensor Product). For a  $k$  tensor  $T$  and a  $j$  tensor  $A$  we define the symmetric tensor product of  $T$  and  $A$  by
$$
T \widetilde {\otimes} A := \operatorname {S y m} (T \otimes A).
$$
Lemma 1. For any tensor  $T$ ,
$$
\| \mathrm {S y m} (T) \| _ {F} \leq \| T \| _ {F}.
$$
$$
\| \operatorname {S y m} (T) \| _ {F} = \left\| \frac {1}{k !} \sum_ {\pi \in S _ {k}} \pi (T) \right\| _ {F} \leq \frac {1}{k !} \sum_ {\pi \in S _ {k}} \| \pi (T) \| _ {F} = \| T \| _ {F}
$$
because permuting the indices of  $T$  does not change the Frobenius norm.
We will use the following two lemmas for tensor moments of the Gaussian distribution and the uniform distribution over the sphere:
Definition 9. For integers  $k, d > 0$ , define the quantity  $\nu_k^{(d)}$  as
$$
\nu_ {k} ^ {(d)} := (2 k - 1)!! \prod_ {j = 0} ^ {k - 1} \frac {1}{d + 2 j} = \Theta (d ^ {- k}).
$$
Note that  $\nu_{k}^{(d)} = \mathbb{E}_{z\sim \mathrm{Unif}(S^{d - 1})}[z_{1}^{2k}]$ .
Lemma 2 (Tensorized Moments).
$$
\mathbb {E} _ {x \sim N (0, I _ {d})} [ x ^ {\otimes 2 k} ] = (2 k - 1)!! I ^ {\widetilde {\otimes} k} \quad a n d \quad \mathbb {E} _ {z \sim \mathrm {U n i f} (S ^ {d - 1})} [ z ^ {\otimes 2 k} ] = \nu_ {k} ^ {(d)} I ^ {\widetilde {\otimes} k}.
$$
Proof. For the Gaussian moment, see [6]. The spherical moment follows from the decomposition  $x \stackrel{d}{=} zr$  where  $r \sim \chi(d)$ .
$$
\| I ^ {\widetilde {\otimes} j} \| _ {F} ^ {2} = \left(\nu_ {j} ^ {(d)}\right) ^ {- 1}
$$
Proof. By Lemma 2 we have
$$
1 = \mathbb {E} _ {z \sim \mathrm {U n i f} (S ^ {d - 1})} [ \| z \| ^ {2 j} ] = \left\langle \mathbb {E} _ {z \sim \mathrm {U n i f} (S ^ {d - 1})} [ z ^ {2 j} ], I ^ {\widetilde {\otimes} j} \right\rangle = \nu_ {j} ^ {(d)} \| I ^ {\widetilde {\otimes} j} \| _ {F} ^ {2}.
$$