Summary
This paper studies the problem of learning single-index models with neural networks. Under a general gradient oracle, the authors present an algorithm that can achieve the optimal computational-statistical tradeoff by obtaining a sample complexity of $\tilde{O}(d^{s^*/2} \lor d)$, and contains the batch-reuse approach of prior work as an example for learning polynomials where $s^* \leq 2$. Furthermore, the authors construct a random loss for each single-index model that implements this general gradient oracle for all $s^*$ when using the $s^*$ degree Hermite polynomial as the activation. Finally, the authors consider the case where the single-index model has additional sparsity, proving an SQ sample complexity lower bound of $n = \tilde{\Omega}(k^{s^*})$ for $k = o(\sqrt{d})$, and presenting an algorithm that matches this lower bound.
Strengths
The paper has solid technical contributions and presents an alternative to the tensor power iteration algorithm of Damian et al., 2024 for learning single-index models with an algorithm that instead relies on a training procedure for neurons. In particular, having a framework that encompasses both batch-reusing and loss modification approaches of prior works is valuable. Further investigating the sparse case presents a nice example of additional hidden structure that can benefit sample complexity.
Weaknesses
* The original motivation as stated in the abstract is to see whether neural networks trained with gradient-based methods can obtain the optimal computational-statistical tradeoff. But the algorithm presented, while interesting from a technical perspective, has little resemblance to the training of neural networks in practice. In that sense, the problem of having a more standard gradient-based algorithm that achieves this tradeoff is still open.
* I believe the writing can be improved in terms of clarity and citing relevant works, with some examples provided below.
Questions
* There are some ambiguities in the discussion starting from Line 258:
* The paragraph before introduces a general class of gradients $\psi(y,x)z$, but here we switch back to activations. What kind of loss function are we using here? The calculation seems to be for squared loss, but my impression was that we have to change the loss to go beyond correlational queries in this online setting.
* The identity in Line 267 might require further clarification. For example, I think $E_Q[y^2]$ should be replaced with $E_Q[\zeta_s(y)y^2]$. Maybe we can apply Cauchy-Schwartz from here to use the assumption on $\zeta_s$, but this would lead to the fourth moment of $y$ appearing. Perhaps the authors could add more details on the derivation to make it clearer.
* My impression was that the only benefit of weight perturbation is to give a smaller denominator in the SNR, however, that doesn't match the following argument. If I try to replicate the SNR calculation of the paragraph starting at Line 220 for this paragraph, the numerator in the SNR is still $d^{-(s^*-1)}$ while the new denominator coming from weight perturbation is $d^{-(s^*-3)/2}$. This leads to an SNR of order $d^{-(s^*+1)/2}$. We argued we need $\sqrt{n\mathrm{SNR}} \gg d^{-1/2}$, therefore $n$ should be of the order $d^{(s^*-1)/2}$. Are there other effects happening when we perturb the weights?
* Line 320 mentions that $L = \tilde{\Omega}(n/\sqrt{d})$ is sufficient. Given $n = d^{s^*/2}\lor d$, this leads to $L = d^{(s^*-1)/2}\lor d^{1/2}$ (which seems to be used in Remark 4.3 as well), which is different from $L$ in Theorem 4.2.
* The loss defined at the end of Page 22 is not exactly a loss function, for example it seems that it can decrease while $y$ is moving away from zero due to the derivative being negative. If this interpretation is correct, perhaps it could be highlighted in the main text that the goal of this loss is not to measure the difference with ground-truth labels, but to construct suitable gradient oracles.
* Can Assumption C.2 be verified for distributions with high generative exponent?
* Theorem 4.2 only states a lower bound on the learning rate. To me it seems like the ideal scenario would be to let the learning rate go to infinity, and directly update the model weights with the gradients as in Algorithm 2. Is there a reason we would want to have a finite learning rate in Algorithm 1?
* Relevant prior works:
* Theorem 5.4 resembles Theorem 3.1 of [VE24], which presents the CSQ lower bound under sparsity, where generative exponent is replaced with information exponent. I think having a discussion on the similarities and differences of the two statements and techniques would be interesting.
* The authors cite [Bach17] as an example of neural networks achieving $O(d)$ sample complexity for learning single-index models. From [Bach17, Table 1], it seems the sample complexity of learning single-index models is $O(d^2)$. In [Bach17, Table 2], the sample complexity will be $d$ with a sign activation and $d^2$ with the ReLU activation. There is a discussion on the sample complexity obtained from [Bach17] in [MHWE24], where the authors also show that standard gradient-based algorithms can achieve the optimal $O(d)$ sample complexity, albeit with a number of queries going beyond the SQ lower bound. This might be interesting to point out while having the discussion on information-theoretically optimal sample complexity with neural networks.
* Another example of additional structure reducing sample complexity is structure in the input covariance, which is known to help in the CSQ setting [BESWW23, MHWSE23]. Developing algorithms that benefit this structure to achieve a smaller sample complexity in the SQ setting can be an interesting direction for future work.
---
References:
[VE24] N. M. Vural and M. A. Erdogdu. "Pruning is Optimal for Learning Sparse Features in High-Dimensions". COLT 2024.
[Bach17] F. Bach. "Breaking the Curse of Dimensionality with Convex Neural Networks". JMLR 2017.
[MHWE24] A. Mousavi-Hosseini, D. Wu, M. A. Erdogdu. "Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics". arXiv 2024.
[BESWW23] J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu. "Learning in the presence of low-dimensional structure: a spiked random matrix perspective". NeurIPS 2023.
[MHWSE23] A. Mousavi-Hosseini, D. Wu, T. Suzuki, M. A. Erdogdu. "Gradient-Based Feature Learning under Structured Data". NeurIPS 2023.