Limitations of SGD for Multi-Index Models Beyond Statistical Queries

Understanding the limitations of gradient methods, and stochastic gradient descent (SGD) in particular, is a central challenge in learning theory. To that end, a commonly used tool is the Statistical Queries (SQ) framework, which studies performance limits of algorithms based on noisy interaction with the data. However, it is known that the formal connection between the SQ framework and SGD is tenuous: Existing results typically rely on adversarial or specially-structured gradient noise that does not reflect the noise in standard SGD, and (as we point out here) can sometimes lead to incorrect predictions. Moreover, many analyses of SGD for challenging problems rely on non-trivial algorithmic modifications, such as restricting the SGD trajectory to the sphere or using very small learning rates. To address these shortcomings, we develop a new, non-SQ framework to study the limitations of standard vanilla SGD, for single-index and multi-index models (namely, when the target function depends on a low-dimensional projection of the inputs). Our results apply to a broad class of settings and architectures, including (potentially deep) neural networks.

💡 Research Summary

The paper tackles a fundamental gap in our theoretical understanding of stochastic gradient descent (SGD). While the statistical query (SQ) framework has become a popular tool for proving lower bounds on learning problems, its assumptions about noise—often adversarial, isotropic, or otherwise artificial—do not faithfully capture the noise generated by standard SGD on real data. Consequently, SQ‑based lower bounds can be either overly pessimistic or simply inapplicable to vanilla SGD, especially when additional algorithmic tricks (sphere constraints, tiny learning rates, batch normalization, etc.) are absent.

To address this, the authors develop a new, non‑SQ analytical framework focused on single‑index and multi‑index models, where the target function depends only on a low‑dimensional linear projection (U x) of the high‑dimensional input (x\in\mathbb{R}^d). Predictors are of the form (f_\theta(x)=h(Wx;\bar\theta)), i.e., any network whose first layer is linear (including deep MLPs, transformers, etc.). The key difficulty is that if the learned subspace (\operatorname{Row}(W)) is poorly aligned with the true subspace (\operatorname{Row}(U)), the gradients carry almost no signal about the target, and the stochastic gradient noise dominates the dynamics.

The authors formalize this intuition through two main constructs:

1. Alignment measure (\rho(W,U)=|P_W P_U|_{\text{op}}), the cosine of the smallest principal angle between the two subspaces. When (\rho) is small, the predictor’s representation is essentially orthogonal to the task‑relevant directions.

2. Gradient condition number (\kappa_T), which quantifies how “well‑behaved” the stochastic gradients are over a horizon of (T) steps. Roughly, (\kappa_T) is the ratio of the worst‑case squared gradient magnitude to the average squared magnitude. If (\kappa_T) does not grow with the ambient dimension (d), the updates of (W) behave like an isotropic random walk.

In the simplest case of a width‑one network ((W) reduces to a vector (w)), the SGD update can be written as (w_T = w_0 - \eta \sum_{t=1}^T a_t x_t), where (a_t) is a scalar depending on the current loss and sample. Under the condition that (\mathbb{E}_t