Two-Point Deterministic Equivalence for Stochastic Gradient Dynamics in Linear Models

Technical Summary: Two-Point Deterministic Equivalence for Stochastic Gradient Dynamics in Linear Models

Problem Statement
Modern deep learning exhibits predictable scaling laws, where performance improves with increasing data volume, model size, and computational capacity. While recent works have analyzed the scaling behavior of kernel methods and random feature models trained with Stochastic Gradient Descent (SGD), existing theoretical frameworks often treat specific effects—such as finite data size, finite model width, and SGD noise—in isolation or rely on complex Dynamic Mean-Field Theory (DMFT) approaches. A unified analytical framework capable of capturing the interplay of limited data, limited features, and SGD noise in high-dimensional linear models, particularly for time-dependent dynamics and Out-of-Distribution (OOD) scenarios, remains a challenge. In particular, standard random matrix methods often rely on "one-point" resolvents, which are insufficient when the matrices governing the dynamics (e.g., empirical covariance, population covariance, and feature matrices) do not commute.

Methodology
The authors develop a novel analytical framework based on two-point deterministic equivalences for resolvents of random matrices. The core of the methodology includes:

1. Model Setup: The article considers linear regression and linear models with random features ($f(x) = x^\top F v$) trained via SGD with a batch size $B$ on a dataset of size $P$. The analysis is conducted in the high-dimensional limit, where the input dimension $D$, feature width $N$, and data size $P$ simultaneously tend to infinity while their ratios remain constant.
2. Reduced Dynamics: Building on previous work, the authors derive a reduced model for the evolution of training and generalization errors. This involves tracking the covariance of the weight deviation ($\Delta w_t \Delta w_t^\top$) and expressing the risk dynamics as a Volterra integral equation. This equation consists of a "driving term" (gradient flow component) and a "kernel term" (SGD noise component).
3. Fourier Space Analysis: The authors transform the time-domain dynamics into Fourier space. This reveals that the driving terms involve products of two resolvents evaluated at different arguments (frequencies), i.e., $(\lambda + AB)^{-1} M (\lambda' + BA)^{-1}$.
4. Two-Point Deterministic Equivalence: The primary technical innovation is the derivation of a deterministic equivalent for these two-point resolvent products. Using the theory of free probability, specifically the properties of the S-transform and a planar diagrammatic expansion, the authors evaluate the asymptotic behavior of these products. They distinguish between "disconnected" contributions (which factorize into one-point equivalents) and "connected" contributions involving mixed cumulants, leading to a self-consistent equation for the two-point function.
5. Application to Risk: These equivalences are applied to derive sharp asymptotic expressions for both training and test risk (generalization), including scenarios with covariate shift (where test data follows a different distribution than the training data).

Main Contributions

• Novel Deterministic Equivalence: The article derives a new class of "two-point deterministic equivalences" for resolvents of random matrices evaluated at arbitrary pairs of arguments. This extends previous "one-point" results to handle non-commuting matrices, which is essential for finite-time dynamics and finite data sizes.
• Unified Derivation: The framework provides a unified derivation of performance for high-dimensional linear regression, kernel regression, and linear models with random features. It recovers known asymptotics (e.g., from DMFT approaches in [16] and one-point equivalence techniques in [17]) while simultaneously delivering new results for time-dependent dynamics.
• Characterization of SGD Noise: The authors explicitly model the interaction of limited data, limited features, and SGD noise. They derive sharp asymptotics for the "SGD kernel term" and demonstrate how noise contributes to risk dynamics over time.
• Analysis of Covariate Shift: The methodology is extended to Out-of-Distribution scenarios, providing exact asymptotic formulas for generalization error when test covariates follow a different Gaussian distribution than the training dataset.
• Connection to DMFT: The article establishes a direct connection between the derived random matrix quantities and Dynamic Mean-Field Theory. It interprets the S-transform of free probability as a response function and identifies the two-point deterministic equivalents with the correlation functions used in DMFT.

Results
The authors obtain explicit, sharp asymptotic formulas for:

• Training and Test Loss: Expressions for the temporal evolution of training and generalization errors in linear models with random features, decomposed into gradient flow driving terms and SGD noise kernel contributions.
• Gradient Flow Terms: Novel two-point equivalents for the driving terms that capture the non-commutative nature of empirical and population covariances at finite times.
• SGD Kernel Terms: Asymptotic expressions for the kernels governing the accumulation of SGD noise, including corrections for finite batch sizes and feature dimensions.
• Static Limits: The results recover known static (infinite time) limits, such as the performance of ridgeless regression, thereby validating the dynamic approach.
• OOD Performance: Closed asymptotic formulas for generalization error under covariate shift, recovering and extending previous static results to the dynamic setting.

Significance
The article claims significance by providing a unified and rigorous analytical tool for understanding the scaling laws and learning dynamics of high-dimensional linear models. By introducing two-point deterministic equivalences, the authors bridge the gap between the static theory of random matrices and the dynamic, non-commutative nature of SGD training. This enables the precise characterization of how finite data, finite model width, and stochastic noise interact to determine performance. The work suggests that the S-transform and planar diagrammatic methods offer a powerful alternative to DMFT for deriving these results, potentially simplifying the analysis of complex learning dynamics while maintaining mathematical rigor in the high-dimensional limit. The authors note that while they do not prove quantitative error bounds for finite dimensions, the asymptotic formulas are expected to provide good approximations based on known results in static settings and numerical evidence from prior literature.