Learning Orthogonal Multi-Index Models: A Fine-Grained Information Exponent Analysis

Thank you for your constructive review and suggestions! We'd like to address your concerns as follows.

Comparison with the generative exponent results

We agree that for single-index models, the information exponent may not be the ideal way to characterize the sample complexity of this task outside the context of one-pass SGD. However, we'd like to emphasize that our results are more about the difference between single- and multi-index models, rather than showing information exponent is not a good way to analyze single-index models. In particular, the argument in our paper can also be generalized to generative exponents (see bellow). We choose to use one-pass SGD and information exponent only because the algorithm and the examples are simpler in this setting.

First, note that even if we replace the information exponent with the generative exponent, we still can't distinguish the cases $\phi = h_2$ and $\phi = h_2 + h_4$. In both cases, the generative exponents are $2$. However, in the first case, due to the rotational invariance, any vector that lies in the relevant subspace is a minimizer of the loss, and it is impossible to recover the exact directions, while in the second case, we can recover the ground-truth directions using the higher-order terms. In other words, in certain multi-index cases, it is still insufficient to consider only the "lowest order", even when the "order" is described using the generative exponent.

On the other hand, we believe that it should also be possible to generalize our positive results to generative exponents. Namely, if the link function is $\phi = h_2 + g$ where $g$ is a reasonable function with generative exponent larger than $2$, then we can recover the exact directions and the sample complexity is still $\tilde{O}(d)$. Similar to the information exponent setting, the intuition here is that we can recover the relevant subspace using $h_2$, and use that as a warm start to learn the exact directions with $g$. (Formally establishing this might be more challenging since the algorithm needed in Stage 2 would be more complicated.)

In summary, our work highlights the difference between singe- and multi-index models, whereas the generative exponent and batch-reusing approaches focus on a different aspect. In particular, both information exponent and generative exponent can fail to capture the richer structure of multi-index models, and in certain cases, this can be fixed by considering the higher order terms.

Responses to other questions/suggestions

• Line 280: No, $k \in [d]$ is not a typo. Here, $v$ refers to an arbitrary first layer neuron and $v_k$ denotes its $k$-th coordinate. We use the indicator $1\{ k \le P \}$ in the formula to differentiate between the cases $k \le P$ and $k > P$.
• Indexing: We will use a different symbol to index the neurons in the revision for clarity.
• Label noise: We believe that our argument still works in the presence of label noise. Similar to situations of the information exponent paper and its follow-ups, the gradient noise is already fairly large, and we have chosen $\eta = \tilde{O}(1/d)$ to make the dynamics stable. Hence, the label noise can usually be merged into the gradient noise as long as it is not too crazy.
• Difference with Ben Arous et al. (2024): The main difference between our results is that in their paper, still only lowest order is considered, while we exploit both the lower and higher order terms. In addition, they run SGD on the Stiefel manifold, which explicitly encodes the non-degeneracy condition (while we need to prove it in Stage 1). The advantage of the Stiefel SGD is that they can handle the case where the second-order terms are not isotropic while we need them to be isotropic.
• Relationship to [7]: Similar to Ben Arous et al. (2024), they run GF on the Stiefel manifold (which prevents collapse) and consider only the lowest order (though the lowest orders can be different for different directions in their setting).

Other related literatures

Thank you very much for pointing out the related literature! We have included them in the revision.

I thank the authors for their rebuttal and for addressing my specific questions.

On the general comment: I understand that the focus of this work is on the difference between single- and multi-index models (and not IE vs. GE).

However, the point I tried to convey is that, in any case, the IE is not a robust notion of complexity. It fails as soon as you change model (single- to multi-index) or algorithm (one pass to full batch / multi-pass SGD), as shown by many works in the literature at this point in time.

Note I am not saying your result is not interesting (I think it is). I am just pointing out that the way you pitch it in the abstract and introduction puts too much weight on the IE as fundamental, while it is already well understood it is a rather restrictive notion. Your result is yet another case example of this.

Thank you for your positive review! We'd like to answer your questions as follows.

Numerical experiments

We have included experimental results in the revision (Appendix G) for link functions $h_2 + h_4$, $h_2 + h_3$ and the absolute value function, with $d$ ranging from $20$ to $200$. For the link function $h_2 + h_4$, we also include experiments where $a$ and $V$ are trained simultaneously. In all cases, the sample/time complexity scales approximately linearly with $d$, which matches our theoretical analysis.

Other questions

• Orthogonal weights: We believe that strict orthogonality is not necessary, but having at least approximately orthogonal weights is important for the argument to work. The approximate orthogonality can be formalized as an incoherent condition, which is also used in Oko et al., 2024 and is also commonly assumed in the tensor decomposition literature. In particular, if the weights are Gaussian vectors, then with high probability, they are approximately orthogonal to each other (and therefore incoherent). In the incoherent case, the dynamics are similar to the orthogonal case, but the proof will be more complicated due to some technical details.
• It is true that when $a_0$ is not small, the dynamics of the weights will no longer be approximately independent. However, in the orthogonal case, the influence of the interaction between neurons can usually be upper bounded. Similar arguments have appeared in, for example, [1] and [2] (though the settings are different). In addition, if we initialize $a_0$ to be small, but do not fix them during training, one can show that $a_k$ will grow faster when $v_k$ is close to one of the ground-truth, as its negative gradient is roughly proportional to $f_*(x) \phi(v_k \cdot x)$. In some sense, the behavior is similar to our layer-wise training strategy. See Appendix G for related experiments. We choose to use layer-wise training only because the proof will be cleaner.
• The only potential issue $h_2 + h_3$ can cause is the initial correlation $v \cdot v^*_k$ might be negative (which is not an issue if the link function is even). This can be easily fixed by replacing each neuron $v$ with two neurons $v$ and $-v$ at initialization. Experimental results supporting this are provided in the Appendix G of the revision.
• The purpose of $l_2$ regularization is mostly technical, as we wish to use the generalization results for ridge regression in Stage 2. We believe that with some effort, one can remove the $l_2$ regularization or replace ridge regression with SGD since after all, the optimization task in Stage 2 (training the second layer with the first layer fixed) is convex.

[1] Yuanzhi Li, Tengyu Ma, and Hongyang R. Zhang. Learning Over-Parametrized Two-Layer Neural Networks beyond NTK. 2020

[2] Gérard Ben Arous, Cédric Gerbelot, and Vanessa Piccolo. High-dimensional optimization for multi-spiked tensor. 2024.

Thank you for your review and suggestions! We'd like to address your concerns as follows.

More comparison with Oko et al., 2024

We wish to emphasize that our result is not a direct extension of Oko et al., 2024 (even when they learn the quadratic part of the link function first). Our sample/time complexity scales linear with $d$ (up to some logarithmic factors), while their results scale with $d^p$. Even if they use the strategy in [DLS22] to learn the quadratic part first and then subtract it from the target function (as they discussed in their paper) to make the information exponent larger than $2$, the final complexity will still be $d^p$ or at least $d^2$, depending on how you implement the algorithm.

The most straightforward implementation of their argument is first learning the quadratic part $f_2$ of $f_*$ and then fitting the function $f_* - f_2$ using their argument. In this case, the complexity is still $d^p$ as this algorithm does not leverage the subspace information learned from the quadratic part.

To utilize the subspace information, they would need to (re)initialize the neurons in the learned subspace (which itself is not very natural) in the second part of the algorithm. In addition, for the one-large-step argument in [DLS22] to work, $d^2$ samples are required. In contrast, our method only requires $d$ samples.

[DLS22] Alex Damian, Jason D. Lee, Mahdi Soltanolkotabi. Neural Networks can Learn Representations with Gradient Descent. 2022

On the scope of this paper

It's true that our results focus on functions with information exponent $2$. However, one goal of our paper is to highlight that information exponent (or generative exponent, see the discussion with reviewer Jcgf) is too coarse to capture the richer structure of multi-index models.

We also believe this type of analysis (starting with lower-order terms and then moving to higher-order terms) can potentially be useful in the general non-orthogonal case. Intuitively, when the order is high, the influence of each ground-truth neuron will be more localized, allowing the nearby neurons to converge to a single direction rather than a mixture of them. However, this localization will also make finding those local attraction regime more difficult. In the ideal case, we should be able to utilize the lower order terms first to find a not-so-local attraction regime and gradually move to higher order terms to localize the learner neurons. This mirrors our approach of using the 2nd order terms to find a right subspace and then the higher order terms to recover the exact directions.

More general link functions

As we have discussed in the paper, we believe that our results can be readily generalized to more general even link functions. In addition, for link functions like $h_2 + h_3$, the only issue here is that the initial $v$ can be negatively aligned with a ground-truth directions. This can be easily fixed by replacing each neuron $v$ with two neurons $v$ and $-v$ at initialization. We include experimental results for link functions $h_2 + h_4$, $h_2 + h_3$ and the absolute value function in the revision (Appendix G). In all cases, the complexity scales roughly linearly with $d$.

Typos and other clarifications

Thank you for identifying the typos! We have corrected them in the revision.

• $\epsilon_*$ and $\epsilon$: In Theorem 2.1, $\epsilon$ and $\epsilon_*$ are the same. This is a typo.
• Normalization step: Indeed, we forget the normalization step in the Eq. (3). However, in the proof, we do have the normalization step (cf. line 1006). For gradient flow, the normalization step is not necessary as long as we project away the radial component of the gradient using $I - vv^T$ (as in Section 3). In GF, if $v_0$ is on the unit sphere, $v_t$ will remain on the unit sphere.
• Introduction of $v_k$: We assume w.l.o.g. that $v^*_k = e_k$ in line 274 and here $v_k$ refers to the $k$th coordinate of $v$.
• Missing $\| v_{\le P} \|^2$ in line 304: Thank you for pointing this out!
• Line 352: This term is correct. It is indeed $c_L(1 - v_1^2)$ instead of $c_L ( \| v_{\le P} \|^2 - v_1^2 )$. We have $v_1^2 \ge (1 + c) v_k^2$ for all $k \ge 2$ and therefore $v_k^2 \lesssim (1 - c) v_1^2$. Plug this into line 345 and then line 340, and we obtain the middle term of line 352. Then, replacing $\| v_{\le P} \|^2$ with $1$ leads to the final inequality.