Research map / Optimization / Implicit bias and generalization

High-dimensional Limit of SGD for Diagonal Linear Networks

NeurIPS 2026 · Optimization › Implicit bias and generalization · approach: Implicit regularization of SGD · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows that streaming SGD on diagonal linear networks in high dimension concentrates around a homogenized SDE and a deterministic PDE for resolvent statistics, with exponential risk decay for small stepsizes under squared parametrization.

Problem

Analyzing SGD with fixed, non-vanishing stepsizes on nonlinearly parametrized models (diagonal linear networks) is hard; classical diffusion approximations need vanishing stepsizes, so an accurate finite-stepsize description of SGD dynamics is lacking.

Approach

In high dimensions, SGD on diagonal linear networks concentrates around a homogenized SDE (drift and noise both proportional to stepsize) and a deterministic integro-differential PDE for a 3x3 resolvent-based statistic; statistics of the form g(resolvent polynomials) inherit this deterministic description, valid for fixed large stepsizes.

Key results

• Theorems give SGD and homogenized SGD concentrating around a deterministic PDE solution with error d^{-epsilon} w.o.p. for fixed stepsizes, including above the convergence threshold (numerical simulations, d=10^3).
• For sufficiently small stepsize under squared parametrization with unit initialization, the risk of the homogenized SDE decays as C e^{-mu t} with high probability.
• Homogenized SGD tracks SGD more accurately than stochastic gradient flow as stepsize grows, whereas SGF degrades.

Surprising findings

• The homogenized SDE remains accurate even for stepsizes above the convergence threshold.

Limitations

• Covariance K must be diagonal for the PDE closure; non-diagonal covariance only supported numerically.
• PDE is nonlinear nonlocal and global well-posedness is not established.
• Exponential decay proof restricted to isotropic squared parametrization with small stepsize; non-explosiveness assumed in general.

Experimental setup

Metrics: risk Hessian trace curvature absolute risk discrepancy vs SGD

Baselines: stochastic gradient flow (SGF) gradient flow

Builds on: stochastic gradient flow high-dimensional deterministic limits of online SGD (ODE characterizations) Gaussian equivalence

Closest papers by meaning

• Statistical Guarantees for High-Dimensional Stochastic Gradient Descent (NeurIPS 2025 · similarity 0.90)
• Learning quadratic neural networks in high dimensions: SGD dynamics and scaling laws (NeurIPS 2025 · similarity 0.89)
• SGD vs GD: Rank Deficiency in Linear Networks (NeurIPS 2024 · similarity 0.89)
• The High Line: Exact Risk and Learning Rate Curves of Stochastic Adaptive Learning Rate Algorithms (NeurIPS 2024 · similarity 0.89)
• The Global Convergence Time of Stochastic Gradient Descent in Non-Convex Landscapes: Sharp Estimates via Large Deviations (ICML 2025 · similarity 0.89)
• Fast Last-Iterate Convergence of SGD in the Smooth Interpolation Regime (NeurIPS 2025 · similarity 0.88)