Interactive Learning of Single-Index Models via Stochastic Gradient Descent

Nived Rajaraman; Yanjun Han

stat.MLFeb 19, 2026

Interactive Learning of Single-Index Models via Stochastic Gradient Descent

Authors: Nived Rajaraman, Yanjun Han

Organizations: Microsoft Research NYC · Courant Institute of Mathematical Sciences and Center for Data Science, New York University

Abstract

Stochastic gradient descent (SGD) is a cornerstone algorithm for high-dimensional optimization, renowned for its empirical successes. Recent theoretical advances have provided a deep understanding of how SGD enables feature learning in high-dimensional nonlinear models, most notably the single-index model with i.i.d. data. In this work, we study the sequential learning problem for single-index models, also known as generalized linear bandits or ridge bandits, where SGD is a simple and natural solution, yet its learning dynamics remain largely unexplored. We show that, similar to the optimal interactive learner, SGD undergoes a distinct burn-in'' phase before entering the learning'' phase in this setting. Moreover, with an appropriately chosen learning rate schedule, a single SGD procedure simultaneously achieves near-optimal (or best-known) sample complexity and regret guarantees across both phases, for a broad class of link functions. Our results demonstrate that SGD remains highly competitive for learning single-index models under adaptive data.

Figures & tables

Panel of Figure 1 — Figure 1: Correlation $m_{t}=\langle\theta_{t},\theta^{\star}\rangle$ plotted as a function of $t$ in $d=20$ dimensions, for the cubic link $f(x)=x^{3}$ . We run interactive SGD with a constant learning rate $\eta_{t}=0.002$ for all $t$ , using an exploration schedule with $\sigma_{t}=0.5$ until $m_{t}$ reaches $0.7$ and $\sigma_{t}=0.2$ thereafter.

Panel of Figure 2 — Figure 2: An example behavior of SGD for pure exploration in the learning phase (cf. Lemma 3.1 ). For appropriately chosen learning rates, if the correlation hits $m_{t}\geq 1-\varepsilon$ at time $t$ , the SGD dynamics will enjoy the following behaviors with high probability: $(i)$ the trajectory will never degrade too significantly, satisfying $m_{s}\geq 1-2\varepsilon$ for all $t\leq s\leq t+\Delta$ ; $(ii)$ at some time $s=T_{0}\in[t,t+\Delta]$ , $m_{s}$ improves to at least $1-\frac{\varepsilon}{4}$ ; and $(iii)$ thereafter, $m_{s}$ may decrease, but will never fall below $1-\frac{\varepsilon}{2}$ for all $T_{0}\leq s\leq t+\Delta$ .

Single-Index Models

Single-index models predict outcomes by applying an unknown nonlinear link to a low-dimensional projection of high-dimensional inputs.

See all 16 papers

High-dimensional learning dynamics of multi-pass Stochastic Gradient Descent in multi-index models
Jan 28, 2026Zhou Fan, Leda WangDynamical Mean-Field Theory Stochastic Gradient Descent
Department of Statistics and Data Science, Yale University
Limitations of SGD for Multi-Index Models Beyond Statistical Queries
Feb 5, 2026Daniel Barzilai, Ohad ShamirSingle-Index Models Statistical Learning Theory
Weizmann Institute of Science · University of Toronto
Optimal Regret for Single Index Bandits
May 10, 2026Devdan Dey, Sujoy Bhore, Avishek GhoshMulti-Armed Bandits Single-Index Models
Department of Computer Science and Engineering, Indian Institute of Technology Bombay.

Jan 28, 2026J/K move · Enter open · S save

High-dimensional learning dynamics of multi-pass Stochastic Gradient Descent in multi-index models

Zhou Fan, Leda Wang

Department of Statistics and Data Science, Yale University

Dynamical Mean-Field Theory Stochastic Gradient Descent High-Dimensional Asymptotics

Open paper AskPDF arXiv

The learning dynamics of Stochastic Gradient Descent (SGD) are believed to play an important role in influencing the generalization properties of trained models. While the precise dynamics of how online (one-pass) SGD optimizes the population risk are now understood in a broad range of high-dimensional problems, a comparable understanding of multi-pass SGD for empirical risk minimization is still largely lacking. In this work, we develop an asymptotically exact characterization of the dynamics of multi-pass SGD for a non-convex multi-index model, in the high-dimensional regime where sample size

n

and data dimension

d

increase proportionally. We consider any sub-linear batch size

κ\asymp n^α

with

α\in [0,1)

and a commensurate ``critical'' scaling of the learning rate, and show that the coordinate-wise dynamics of SGD admit a novel continuous-time scaling limit, given by a system of dynamical mean-field equations that are driven by a scalar Poisson jump process representing the asymptotic limit of SGD sampling noise. This scaling limit connects several previously disparate theories around SGD, namely (i) the classical interpretation of SGD as an approximation to gradient flow on the empirical risk, (ii) the characterization of online SGD via an autonomous ODE in its overlap statistics, and (iii) the coincidence between multi-pass SGD dynamics and diffusion approximations in least-squares problems. We illustrate how these perspectives are recovered by specializing our result to small learning rates, large sample sizes, and linear models, respectively. Outside of these specializations, the scaling limit of SGD is distinct both from diffusion approximations and gradient flow, and may enable SGD to converge to local optima of the non-convex empirical risk with different statistical properties.

Abstract

Single-Index Models

Explore similar work

High-dimensional learning dynamics of multi-pass Stochastic Gradient Descent in multi-index models

Limitations of SGD for Multi-Index Models Beyond Statistical Queries

Optimal Regret for Single Index Bandits