Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics

We thank the reviewer for their valuable evaluation and insightful feedback, and address their comments and questions below.

What conditions can lead to non-exponential complexity?

A main challenge in proving polynomial-time convergence here is to find settings in which the log-Sobolev constant scales polynomially with the inverse temperature. Unfortunately, as long as the Gibbs potential has disjoint local minima with large energy barriers, one cannot hope to establish polynomial log-Sobolev inequalities in general [1]. We believe that further analysis on the Gibbs potential can identify settings in which no such energy barrier exists, which would likely include simpler functions such as those with information exponent equal to 1, indeed partial progress towards this has been achieved in a two-timescales setting [2]. This is an interesting direction for future research, as it would provide a unified analysis for learning these simpler functions as opposed to the ad-hoc analyses in the current literature, and would cover a much broader class of input distributions.

The analysis is restricted to 2-layer neural networks and $k \ll d$:

We agree with the reviewer that the setting of our theoretical study is quite simple in comparison with deep neural networks used in practice. However, we remark that two-layer neural networks have been the focus of most recent theoretical work on feature learning with neural networks. There have been attempts at generalizing the mean-field formulation to more than 2 layers, e.g. [3], but much less is known in that setting, and it is an interesting direction to prove learnability guarantees with gradient-based algorithms for deeper networks.
Please note that since the minimax lower bounds in this setting suggest that learning the multi-index model without additional assumptions (such as ridge separability [5]) and with $\epsilon$ error requires $n = \Omega\big((1/\epsilon)^{\Omega(k)}\big)$ samples, a statistically efficient algorithm can only exist when $k \ll d$. However, we believe Proposition 8 provides one interesting example of a setting in which $k$ is large, i.e. the target depends on many neurons, yet both sample and time complexity scale mildly with problem parameters as shown in Theorem 10.

Proving LSI without regularization/with dissipativity:

There are two key challenges in removing the weight decay regularization in our theoretical analysis.
First, as pointed out by the reviewer, this regularization allows us to prove LSI for the Gibbs potential by considering it as a bounded perturbation of a strongly convex potential. We remark that dissipativity-type conditions are not likely to work here. In particular, the gradient of the Gibbs potential with respect to the weights will have a bounded norm (for activations with bounded first derivative), i.e. $\Vert \nabla_w \hat{\mathcal{J}}_0’ [\mu] (w) \Vert \leq C$ for all $w$, thus $\langle w, \nabla_w \hat{\mathcal{J}}_0’ [\mu] (w) \rangle \leq C\Vert w \Vert$, which rules out the possibility of a dissipativity condition.
Second, we also require weight decay to be able to apply our generalization bounds. While in practice it is believed that gradient descent will have an inductive bias towards minimum norm solutions that generalize well, this phenomenon is not well studied in the setting we consider here. We believe that removing regularization while still showing that gradient-based algorithms learn networks with good generalization properties is an important open problem.

The meaning of $\mathbb{E}_\zeta[\rho(\zeta)]$:

$\rho$ denotes the loss function, $\zeta$ is the additive noise in the multi-index model. This term can be thought of as an irreducible loss in regression.

We thank the reviewer for their valuable evaluation and insightful feedback, and address their comments and questions below.

The comparison with information/generative exponent analysis:

We thank the reviewer for raising their concern about this comparison. We do not claim that MFLA achieves a sample complexity beyond that suggested by computational lower bounds (dictated by the information/generative exponent). Indeed, with exponential width, one may conclude that a small fraction of neurons enter the descent phase upon random initialization. However, as this subset is exponentially small, it is not clear that standard end-to-end training can achieve sufficient amplification of these “warm” neurons after $O(d)$ steps to lower the loss – this will not be the case if we run the spherical gradient algorithm (which is the update studied in most information/generative exponent analyses), because the neural network output is still dominated by the majority of neurons that have not “escaped from mediocrity”.
Moreover, we highlight that the information/generative exponent analysis is only applicable to the Gaussian input distribution setting (with the exception of [1] which only considers rotationally invariant input distributions and known link function), and is further limited in either assuming a known single-index link function and choosing the same link function as the activation, or having nontrivial modifications of the standard training procedure for 2-layer neural networks. Moreover, most results so far capture learnability for single-index models, and characterizing learnability of general multi-index models with this type of analysis is much more challenging, and these analyses typically fall short of providing a concrete sample complexity for such general settings.
Part of our motivation here is to instead provide a general and unifying theory for all multi-index models in a distribution free sense, which comes at the cost of a larger worst-case computational complexity. Moreover, the adaptivity of such analysis to low effective dimension, i.e. low-dimensional input embedded in a high-dimensional space, has not been established in prior works. We hope that this discussion helps provide additional context in comparing our work with the literature.

Examples of natural functions in Proposition 8:

Thank you for the suggestion. We have added an example in Appendix D that we believe will in fact provide further insights to the readers. In summary, we let $\mathcal{W} = \mathbb{S}^{d-1}$ and $\Psi(x;w) = \langle x, w\rangle^2$. We let $\mu^*$ be the distribution of $\frac{A^{1/2} u}{\Vert A^{1/2} u\Vert }$, where $u \sim \mathrm{Unif}(\mathbb{S}^{d-1})$. Then, $\hat{y}(x;\mu^*) = \frac{1}{d}\Vert B^{1/2} x\Vert^2$ where $B = A^{1/2}\mathbb{E}\left[\frac{uu^\top}{\Vert A^{1/2}u\Vert^2}\right]A^{1/2}$, whereas $\hat{y}(x;\tau) = \frac{1}{d}\Vert x \Vert^2$ (please note that we assume the typical norm of $x$ to be of order $\sqrt{d}$). If we simply let $A = \mathrm{diag}(\lambda_1,1,...,1)$ and $\lambda_1 = 1/d$, then $\Vert B - I_d\Vert_{\mathrm{op}} \geq \Omega(1)$. This means that for some distribution over $x$, the initialization $\tau$ has $\Omega(1)$ loss, and one needs to train with MFLD to recover $\mu^*$ and achieve small loss. Furthermore, $\bar{\Delta} \leq \mathcal{O}(\ln d)$ in this setting, thus we can achieve the polynomial convergence guarantee of Theorem 10.

The effective dimension does not capture the latent low dimension in the multi-index model:

We would like to highlight that there are two sources of potential low-dimensionality, and only having a small number of indices $k$ is not sufficient to learn with linear sample complexity while having polynomial time convergence, as evident by computational lower bounds [2,3]. While for isotropic input covariance we can achieve the optimal linear sample complexity regardless of information/generative exponent, we believe that the additional benefit in sample and computational complexity thanks to the low effective dimension is an interesting phenomenon which is mostly overlooked in the current literature on feature learning of neural networks, especially since in practice directions of high variation tend to be better predictors of the label, i.e. many problems tend to have a low effective dimension [4].

We would be happy to clarify any concerns or answer any questions that may come up during the discussion period.

References:

[1] A. Zweig et al. “On single index models beyond gaussian data.” NeurIPS 2023.

[2] A. Damian et al. "Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models." NeurIPS 2023.

[3] A. Damian et al. "Computational-Statistical Gaps in Gaussian Single-Index Models." COLT 2024.

[4] T. Hastie et al. "The elements of statistical learning: data mining, inference, and prediction." 2009.

We thank the reviewer for their valuable evaluation and insightful feedback, and address their comments and questions below.

• A discussion on when polynomial-time convergence is achievable:

A main challenge in proving polynomial-time convergence here is to find settings in which the log-Sobolev constant scales polynomially with the inverse temperature. Unfortunately, as long as the Gibbs potential has disjoint local minima with large energy barriers, one cannot hope to establish polynomial log-Sobolev inequalities in general [1]. We believe that further analysis on the Gibbs potential can identify settings in which no such energy barrier exists, which would likely include simpler functions such as those with information exponent equal to 1, indeed partial progress towards this has been achieved in a two-timescales setting [2]. This is an interesting direction for future research, as it would provide a unified analysis for learning these simpler functions as opposed to the ad-hoc analyses in the current literature, and would cover a much broader class of input distributions.

We would also like to highlight that we have added a natural example of a target function in the form of Proposition 8 in Appendix D.

We would be happy to clarify any concerns or answer any questions that may come up during the discussion period.

References:

[1] G. Menz and A. Schichting. “Poincaré and logarithmic Sobolev inequalities by decomposition of the energy landscape.” Annals of Probability, 2014.

[2] G. Wang et al. “Mean-Field Langevin Dynamics for Signed Measures via a Bilevel Approach.” arXiv 2024.

We thank the reviewer for their valuable evaluation and insightful feedback, and address their comments and questions below.

On Brownian noise helping to escape from local minima:

Indeed, here we add artificial noise different from noise arising from mini-batch sampling in SGD. This noise is crucial for convergence analysis, and intuitively allows the dynamics to escape from saddle points and local minima. Please see [1, Section 5.2] for an example of this phenomenon with mean-field Langevin dynamics, or [2] and references therein for using additional noise in GD/SGD to escape saddle points. We will further highlight this discussion in the final version.

The generality of Assumption 1:

We thank the reviewer for raising their concern here. Our assumption here is much broader than the Gaussian data $x \sim \mathcal{N}(0,I_d)$ assumption of the recent literature on feature learning of neural networks, e.g. [4, 5, 6]. Indeed, Assumption 1 is automatically verified when $x \sim \mathcal{N}(0,\Sigma)$, with $\sigma_n = \sigma_u = \mathcal{O}(1)$. We would like to highlight however that Assumption 1 is much broader than the Gaussian case. For example, if $x = \Sigma^{1/2} z$ where $z$ has i.i.d. mean-zero subGaussian coordinates, then Assumption 1 is verified with $\sigma_u = \sigma_n = \mathcal{O}(1)$ [3, Theorem 6.3.2]. We have clarified this point in the revised version.

What is $\mathbb{E}_\zeta[\rho(\zeta)]$ in Theorem 3?

$\rho$ denotes the loss function, $\zeta$ is the additive noise in the multi-index model. This term can be thought of as an irreducible loss in regression.

Numerical Simulation:

Thank you for this suggestion. We have added a simple numerical simulation to verify the intuitions from our theory, specifically that the effective dimension controls the generalization gap, in Appendix E. We will provide more extensive numerical simulations in the camera-ready version.

Further, we thank the reviewer for mentioning the typo between lines 230 and 231, which we have fixed.

We would be happy to clarify any concerns or answer any questions that may come up during the discussion period.

References:

[1] L. Chizat. “Mean-field langevin dynamics: Exponential convergence and annealing.” TMLR 2022.

[2] H. Daneshmand et al. “Escaping Saddles with Stochastic Gradients.” ICML 2018.

[3] R. Vershynin. “High-Dimensional Probability: An Introduction with Applications in Data Science.”

[4] A. Damian et al. “Neural networks can learn representations with gradient descent.” COLT 2022.

[5] J. Ba et al. “High-dimensional asymptotics of feature learning: how one gradient step improves the representation.” NeurIPS 2022.

[6] A. Bietti et al. “Learning Single-Index Models with Shallow Neural Networks.” NeurIPS 2022.