Exact Dynamics of Multi-class Stochastic Gradient Descent

TL;DR

This work analyzes the exact high-dimensional dynamics of online stochastic gradient descent (SGD) for learning with Gaussian mixture data across multiple anisotropic classes. By deriving a deterministic system of integro-differential equations and an autonomous ODE-based framework, the authors characterize the evolution of the iterate norm and the overlap with class means, enabling precise predictions of the population loss and related statistics in the proportional regime where dimension scales with samples. They extend prior single-class results to multi-class settings with a growing number of classes and analyze three covariance structures—identity, zero-one, and power-law—revealing a structural phase transition in which SGD concentrates on directions of smaller variance. The framework also encompasses Gaussian mixture data with mean and covariance heterogeneity and provides a homogenized SGD interpretation, with explicit results for binary logistic regression and a multi-class MSE example. Theoretical predictions are validated by simulations, highlighting how anisotropy drives learning dynamics and potentially accelerates or slows convergence depending on spectral properties. The results offer quantitative insights into SGD behavior on structured, high-dimensional data and contribute to understanding neural network training dynamics through tractable Gaussian toy models.

Abstract

We develop a framework for analyzing the training and learning rate dynamics on a variety of high- dimensional optimization problems trained using one-pass stochastic gradient descent (SGD) with data generated from multiple anisotropic classes. We give exact expressions for a large class of functions of the limiting dynamics, including the risk and the overlap with the true signal, in terms of a deterministic solution to a system of ODEs. We extend the existing theory of high-dimensional SGD dynamics to Gaussian-mixture data and a large (growing with the parameter size) number of classes. We then investigate in detail the effect of the anisotropic structure of the covariance of the data in the problems of binary logistic regression and least square loss. We study three cases: isotropic covariances, data covariance matrices with a large fraction of zero eigenvalues (denoted as the zero-one model), and covariance matrices with spectra following a power-law distribution. We show that there exists a structural phase transition. In particular, we demonstrate that, for the zero-one model and the power-law model with sufficiently large power, SGD tends to align more closely with values of the class mean that are projected onto the "clean directions" (i.e., directions of smaller variance). This is supported by both numerical simulations and analytical studies, which show the exact asymptotic behavior of the loss in the high-dimensional limit.

Figures (11)

• Figure 1: Population risk concentration on a binary logistic regression problem for different values of learning rate $\gamma,$ and for (a) identity covariance with standard normal mean. (b) zero-one model with mean non-zero only in the subspace for which both covariance has zero eigenvalues (c)/(d) power-law diagonal covariance matrices i.e., $\lambda_\rho^{(1)}=(\frac{\rho}{d})^{1.1}$ and $\lambda_\rho^{(2)}=(\frac{\rho}{d})^{1.5}$ and mean $\mu_\rho = \frac{1}{\sqrt{d}}(\frac{\rho}{d})^{\beta}$ for two different values of $\beta = 0.2, 1$.
• Figure 2: Identity and zero-one model with $d=1000$ and $\gamma=0.9$. We display $\mathscr{L}(t)$, $\mathfrak{m}(t)$, $\mathscr{V}(t)$ and the ratio $\mathfrak{m}(t)/\sqrt{\mathscr{V}}$, i.e. the normalized version of $\mathfrak{m}(t)$. On each graph, we plot both the SGD simulation and the theoretical prediction (ODE).
• Figure 3: Zero-one model with $d=1000$ and $\gamma=0.9$ showing subspace alignment
• Figure 4: Power-law model with $d=1000$, with $a_1=a_2=1.2$ and different $\beta$
• Figure 6: Power-law model with $d=1000$, with $a_1=a_2=1.2$ and different $\beta$. Here we project $\mathfrak{m}(t)$ into subspaces associated with the lower half of the eigenvalues and the upper half of the eigenvalues and we refer to these projections as $m_L(t)$ and $m_U(t)$ respectively.

Theorems & Definitions (71)

• Theorem 1: Learning curves
• Remark 1
• Proposition 2
• Remark 2
• Remark 3
• Remark 4: The role of $\alpha$ and $\beta$
• Proposition 3: Mild power-law regime and identity
• Proposition 4: Extreme power-law
• Remark 5
• Lemma 1: Linear least square GMM