Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit

We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.

"Can the use of momentum be avoided in the first layer of training?"

We use an interpolation step to improve the signal-to-noise ratio in the gradient; this is crucial in the generative exponent 2 setting where the signal and noise are almost at the same magnitude -- please see our response to Reviewer p2Vt for more details and a recap of our explanation in Section 3.2.
It is possible that other modifications or hyperparameter choices of the learning algorithm also achieve a similar effect, but we do not pursue these in the current submission.

[Gla'23] and the possibility of going beyond SQ

Thank you for the close reading. We realize that this remark on future direction is misleading. We initially thought that for the $k$-parity problem, the SQ lower bound suggests that $n\gtrsim d^k$ samples are required for rotationally invariant algorithms with polynomial compute (as opposed to $n\gtrsim d^{k/2}$ in the Gaussian case). However this is not the case, and the total computation in [Glasgow 23] only matches that suggested by the SQ lower bound $q/\tau^2\asymp d^2$. We will include appropriate references on the gap between statistical and adversarial noise, such as [Dudeja and Hsu 20] [Abbe and Sandon 20] (albeit with highly nonstandard architecture).
We also agree with the reviewer that the statement needs to be reworded in the absence of empirical evidence.

We would be happy to clarify any further concerns/questions in the discussion period.

We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.

Standard SGD practices, relaxing assumptions

We agree that Algorithm 1 deviates from the most standard training procedure in practice. Note that layer-wise training and fixed bias units are fairly standard algorithmic modifications when proving end-to-end sample complexity guarantees, as seen in many prior works on feature learning [DLS22] [AAM23]. Without such modifications, stronger target assumptions are generally needed, such as matching link function [BAGJ21]. Below we discuss the possible extensions to more general settings.

• Simultaneous training. We believe that the statistical efficiency of Algorithm 1 can be achieved by simultaneous training of all parameters, under appropriate choice of hyperparameters. For instance, the layer-wise procedure may be approximated by a two-timescale dynamics [BBPV23], and the required ``diversity'' of bias units may be met when the learning rate of the bias weights is sufficiently small. As mentioned in the Conclusion section, it is also interesting to consider a more standard multi-pass algorithm instead of the currently employed data-reuse schedule.

• Different losses. We focus on the squared loss as it is the most standard objective for regression tasks. Note that the restriction to correlational information is a feature of the squared/correlation loss -- see Section 2.2. If we employ a different loss such as L1 loss, it is possible that SGD can implement non-correlational queries just from the loss function itself -- such direction is tangential to our analysis, since our goal is to show that such non-CSQ component naturally arises from reuse of training data.

Number of neurons and dependence on target degree $q$

The required width of the neural network is almost dimension-free. In particular, we only need to set $N=\text{polylog}(d)$ to achieve $o_d(1)$ population error -- see line 245 for details.

On the other hand, our big-$O$ notation hides a constant that might depend exponentially on the target degree $q$. This is due to the sign of Hermite coefficients required in Assumption 3, and the compactness argument to uniformly upper-bound $C_q$ in Proposition 4. Note that similar dependence on $q$ is also present in tailored algorithms for learning low-dimensional polynomials [CM20].

"why the even polynomials are harder than the odd ones"

Intuitively speaking, the neural network is initialized at an (approximate) saddle point when $f_*$ is even, since the expectation $\mathbb{E}[\mathcal{T}(y)\langle\boldsymbol{x},\boldsymbol{\theta}\rangle]=0$ for any $\mathcal{T}\in L^2(y)$. See [BAGJ21] for more discussions.

To elaborate, let us consider one-pass SGD and evaluate the scale of population gradient of one neuron $\sigma(\langle \boldsymbol{w},\boldsymbol{x}\rangle)$ at random initialization. Let $\sigma_*=\sum_{i=0}^q c_i He_i$, we have $\mathbb{E}[\nabla_{\boldsymbol{w}}\sigma_*(\langle \boldsymbol{\theta},\boldsymbol{x}\rangle)\sigma(\langle \boldsymbol{w},\boldsymbol{x}\rangle)] \approx \sum_{i=1}^q \mathbb{E}_{t\sim \mathcal{N}(0,1)}[\sigma^{(i)}(t)] c_i\langle \boldsymbol{\theta},\boldsymbol{w}\rangle^{i-1}\boldsymbol{\theta}.$ Therefore, when $c_1 \ne 0$, the scale of the population gradient is $\Theta(1)$, while when $c_1=0$ and $c_2\ne 0$, it is $\langle \boldsymbol{\theta},\boldsymbol{w}\rangle \simeq d^{-\frac12}$ with high probability.

Similarly for reuse-batch SGD, the population gradient is a linear combination of $\mathbb{E}[\nabla_{\boldsymbol{w}}\sigma_*(\langle \boldsymbol{\theta},\boldsymbol{x}\rangle)^i\sigma^{(i-1)}(\langle \boldsymbol{w},\boldsymbol{x}\rangle)(\sigma^{(1)}(\langle \boldsymbol{w},\boldsymbol{x}\rangle))^{i-1}]$. When $\sigma_*$ is even, one can similarly see that the scale of each is $O(d^{-\frac12})$.

"what could be an optimal mini-batch re-use schedule?"

We intuitively expect that any mini-batch size $b$ between $1$ to $d$ with $\eta=\tilde{\Theta}( d^{-1}b)$ would achieve similar sample complexity. But for simplicity, we employed $b=1$, which excludes the correlation between different samples.
Investigating the benefit of more intricate mini-batch schedule is an interesting direction for future work.

We would be happy to clarify any further concerns/questions in the discussion period.

We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.

Polynomial activation function

We make the following clarifications.

1. Note that all square-integrable activation functions (ReLU, sigmoid, etc.) can be written as a linear combination of Hermite polynomials. The key differences between our algorithm and [CM20] are as follows,

• [CM20] employed a label transformation (based on thresholding) prior to running SGD, whereas we use the squared loss without preprocessing and extract the transformation from reusing training data.

• [CM20] considered optimization jointly over the low-dimensional subspace (finding index features) and the coefficients of the polynomials. In contrast, in our setting the coefficients of the polynomial activation are fixed, and we optimize the parameters of the neural network.

2. We restrict ourselves to polynomial activation because it is easy to construct coefficients that satisfy Assumption 3 (required for strong recovery). For strong recovery, ReLU or sigmoid may not suffice, and the use of well-specified or polynomial activations is common in the literature, e.g., [AGJ21][AAM22][AAM23][DNG+23].
   As for weak recovery, Assumption 2 is satisfied with probability 1 by a shifted ReLU/sigmoid (e.g., see Lemma 15 in [BES+23]). Therefore, we can establish weak recovery using standard choices of activation. We will comment on this in the revised manuscript.

Dimension dependence in constants

All constants in our theorems are dimension free. Specifically, the lower bound on $H_0$ does not depend on dimensionality, since Proposition 4 is for univariate functions $f:\mathbb{R}\to\mathbb{R}$. We notice that there is a typo in Appendix A: the expectation in A.1 should be with respect to $\mathcal{N}(0,1)$ since $f$ is a scalar function. We apologize for the confusion this may cause.

"does it imply that the width of the network is at least $\text{exp}(q)$?"

Yes, the required student width is exponential in the target degree $q$, which is treated as a constant the big-$O$ notation. Although only a small fraction of neurons can achieve strong recovery, the entire neural network can achieve small $L^2$ error because these ``good'' neurons can be singled out in the second-layer training, as shown in Proposition 3.
Note that similar dependence on $q$ is also present in tailored algorithms for learning low-dimensional polynomials [CM20].

We would be happy to clarify any further concerns/questions in the discussion period.

Please engage with the authors' response: have they addressed your concerns adequately, and how has your score been affected?

Best, your AC

I thank the authors for their detailed response. I would like to keep my score unchanged.

We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.

"The idea of label transformation implemented by [CM20] could be reported"

In the current manuscript, the difference between our label transformation and that in [CM20] is discussed in the paragraph starting line 157 and line 250. Specifically, prior SQ algorithms are typically based on thresholding, whereas we use monomial transformations which can be easily extracted from SGD update.

"one might think that sample complexity guarantees might come easily from weak recovery plus [BAGJ21]"

The difficulty is discussed in the paragraph starting line 162. In particular, since the link function $\sigma_*$ is unknown, we cannot directly employ the argument in [BAGJ21] which assumed a well-specified model, i.e. $\sigma_* = \sigma$. Instead, we make use of the randomized Hermite coefficients of the student activation function to translate weak recovery to strong recovery.

"What is the role of the hyperparameters ... and all the quantities appearing in the non-CSQ transformation?"

Most hyperparameters in Algorithm 1 are deterministic, with the only exceptions being the sign of student activation function (see Lemma 2) and the momentum parameter $\xi$ randomized over student neurons. This aims to guarantee that, for any target function, there exists some student neurons achieving strong recovery (Theorem 2). We realize that there is a typo in Theorem 1: "with probability" should be replaced by "with high probability". We apologize for the confusion this may cause.
We will clarify this in the revision.

The role and necessity of interpolation

We use an interpolation step to improve the signal-to-noise ratio in the gradient; this is crucial in the generative exponent 2 setting where the signal and noise are almost at the same magnitude (see Section 3.2 and paragraph starting line 307 for details). It is possible that other modifications or hyperparameter choices of the learning algorithm also achieve a similar effect, but we do not pursue these in the current submission.

Below we recap the intuition provided in Section 3.2 on the failure of the standard online SGD analysis. When we analyze the training dynamics of learning single-index model, we characterize the progress via the projection of $\boldsymbol{w}^t$ onto $\boldsymbol{\theta}$, which we refer to as the alignment $\kappa^t = \langle \boldsymbol{\theta}, \boldsymbol{w}^t \rangle$ (e.g., [DNG+23]). We want to show that $\mathbb{E}[\kappa^{t+1}]$ increases from $\kappa^t$. Given the gradient $\boldsymbol{g}^t$ (assumed to be orthogonal to $\boldsymbol{w}^t$ for simplicity) and step size $\eta$, the update of alignment is given as $\kappa^{t+1} = \left\langle \boldsymbol{\theta}, \frac{\boldsymbol{w}^t + \eta \boldsymbol{g}^t}{||\boldsymbol{w}^t + \eta \boldsymbol{g}^t||} \right\rangle \gtrsim \underbrace{\left\langle \boldsymbol{\theta}, \boldsymbol{w}^t \right\rangle}_{=\kappa^t}+ \eta \left\langle \boldsymbol{\theta}, \boldsymbol{g}^t \right\rangle - \frac{1}{2} \eta^2 ||\boldsymbol{g}^t||^2 \kappa^t + \text{(mean-zero noise)}.$

One sees that for the expectation to be larger than $\kappa^t$, $\eta \mathbb{E}[\langle\boldsymbol{\theta}, \boldsymbol{g}^t\rangle]$ should be larger than $\frac12\eta^2 ||\boldsymbol{g}^t||^2\kappa^t$. To achieve this, we can simply let $\eta^t$ sufficiently small ($\Theta(d^{-1})$ for $\mathrm{IE}=2$); because the signal term linearly depends on $\eta$, while the noise term quadratically.

However, in our case, the signal comes from non-CSQ term. For example, when $\mathrm{IE}(y^I)=2$, the signal term is proportional to $(\eta^t)^Id^{I-1}\kappa^t$ (under $\eta^t \lesssim d^{-1}$). Therefore, decreasing $\eta^t$ does not improve the SNR. The interpolation step provides a remedy by preventing the parameters from changing too fast and reducing the projection error.

"it would be great to state more clearly details over the SQ class"

Thank you for the valuable suggestion. We will discuss the SQ complexity and generative exponent in [DPVLB24] in more details. Indeed our analysis is based on the observation that SGD with reused batch can implement SQ, and hence $\tilde{O}(d)$ samples are sufficient to learn target functions with generative exponent at most 2; this sample complexity is also achieved by AMP algorithms (typically assuming knowledge of the link function to construct the optimal preprocessing), and is consistent with the SQ lower bound.

"could the author comment more on the link with [Gla23] and non-adversarial noise and the possibility of going beyond SQ?"

Thank you for the close reading. We realize that this remark on future direction is misleading. We initially thought that for the $k$-parity problem, the SQ lower bound suggests that $n\gtrsim d^k$ samples are required for rotationally invariant algorithms with polynomial compute (as opposed to $n\gtrsim d^{k/2}$ in the Gaussian case). However this is not the case, and the total computation in [Glasgow 23] only matches that suggested by the SQ lower bound $q/\tau^2\asymp d^2$. We will include appropriate references on the gap between statistical and adversarial noise, such as [Dudeja and Hsu 20] [Abbe and Sandon 20].

We would be happy to clarify any further concerns/questions in the discussion period.