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

TL;DR

This work studies learning multi-index models of the form $f_*(x)=\sum_{k=1}^P \phi(\mathbf{v}_k^*\cdot x)$ under Gaussian inputs, where the Hermite expansion of $\phi$ includes a second-order term and a higher-order term at order $L>2$. It shows that vanilla information-exponent analyses can be suboptimal in multi-index settings by exploiting both second-order and higher-order terms: first learning the relevant subspace via second-order information, then recovering the exact directions via higher-order information. The authors present a two-stage online SGD procedure (followed by ridge regression) that achieves a total sample complexity of $\tilde{O}( d P^{L-1} )$, with a supporting gradient-flow analysis and a suite of stochastic-noise lemmas to bridge population dynamics to online learning. This result highlights how higher-order structure can dramatically improve learning efficiency in high-dimensional, low-rank settings and provides analytic tools (stochastic Gronwall, subweibull bounds) for noisy online learning dynamics.

Abstract

The information exponent ([BAGJ21]) and its extensions -- which are equivalent to the lowest degree in the Hermite expansion of the link function (after a potential label transform) for Gaussian single-index models -- have played an important role in predicting the sample complexity of online stochastic gradient descent (SGD) in various learning tasks. In this work, we demonstrate that, for multi-index models, focusing solely on the lowest degree can miss key structural details of the model and result in suboptimal rates. Specifically, we consider the task of learning target functions of form $f_*(\mathbf{x}) = \sum_{k=1}^{P} φ(\mathbf{v}_k^* \cdot \mathbf{x})$, where $P \ll d$, the ground-truth directions $\{ \mathbf{v}_k^* \}_{k=1}^P$ are orthonormal, and the information exponent of $φ$ is $L$. Based on the theory of information exponent, when $L = 2$, only the relevant subspace (not the exact directions) can be recovered due to the rotational invariance of the second-order terms, and when $L > 2$, recovering the directions using online SGD require $\tilde{O}(P d^{L-1})$ samples. In this work, we show that by considering both second- and higher-order terms, we can first learn the relevant space using the second-order terms, and then the exact directions using the higher-order terms, and the overall sample and complexity of online SGD is $\tilde{O}( d P^{L-1} )$.

Figures (1)

• Figure 1: Recovery of directions. The above plots show the evolution of the correlation with each of the ground-truth directions. We fix the relevant dimension $P = 5$ and vary the ambient dimension $d$. Different colors represent different $d$. For each color, one curve represents $\max_{{\bm{v}}} v_p^2$ for one $p \in [P]$. In the first row, the link function is $\phi = h_2 + h_4$. In the left plot, we use the algorithm \ref{['eq: training algorithm']}, while in the right plot, we train both layers simultaneously. The second row contains simulation results for other link functions.

Theorems & Definitions (79)

• Theorem 1.1: Informal version of Theorem \ref{['thm: main']}
• Lemma 2.1: First-layer gradients
• proof : Remark
• Theorem 2.1: Main Theorem
• proof : Remark
• Lemma 4.1: Stochastic Gronwall's lemma
• proof : Condition \ref{['eq: conditions of stochastic gronwall']}
• proof : Stochastic induction
• proof : Remark on the subweibull condition
• proof : Proof sketch of Lemma \ref{['lemma: stochastic discrete gronwall']}