Open problem MAIS-O76 · posed in MAIS-A7 as Problem 3.8 · Status: open.
Tags: generalization · training dynamics · singular learning theory · simplicity bias · probability · algebraic geometry. Difficulty: ★★★.
Authored by: Claude Fable 5 directed by Lionel Levine · Audited by: GPT 5.6 Sol.
What does stochastic gradient descent do after it reaches zero loss? For an overparametrized model the zero-loss parameters form a whole set, typically singular, and which point on it training selects is the implicit bias that determines what the trained model generalizes to. Over a smooth manifold of minimizers the answer is a theorem: Li, Wang, and Arora [LWA22], adapting a 1991 theorem of Katzenberger [K91], prove that SGD with learning rate $\eta \to 0$, suitably rescaled in time, converges to a diffusion on the manifold; under label noise (fresh independent noise added to the regression targets at each step, which makes the noise covariance proportional to the Hessian near the zero set) the limiting drift descends the Hessian trace, so SGD slides along the optimal set toward flatter points, a preference established for label-noise SGD by Damian, Ma, and Lee [DML21].
The theorem's hypothesis — Hessian rank constant along the manifold and equal to its codimension — excludes precisely the singularities where the strata of the optimal set meet, which is where the selection happens. The problem extends the limit theorem across them, with matrix factorization as the concrete instance. In the statement, the Li–Wang–Arora coefficients of a smooth stratum of minimizers are the drift and diffusion coefficients of this limiting diffusion on it, determined by the loss Hessian and the gradient-noise covariance (exact formulas in MAIS-A7, Section 3.4, and in [LWA22]); and $W_0 = L^{-1}(0)$ denotes the optimal set of a loss $L$.
Problem (MAIS-A7, Problem 3.8; effective diffusion beyond the smooth part). Let $L(A, B) = \tfrac12 \Vert BA - \Phi \Vert _F^2$ with $A \in \mathbb{R}^{H \times M}$, $B \in \mathbb{R}^{N \times H}$, $H > r = \mathrm{rank}\ \Phi$, and $H \le \min(M, N)$. On the fiber $F_\Phi = \lbrace (A,B) : BA = \Phi \rbrace $, Sylvester's rank inequality $\mathrm{rank}\ A + \mathrm{rank}\ B - H \le \mathrm{rank}\ BA = r$ caps the rank sum, so the top stratum is
$$\Gamma = \lbrace (A,B) \in F_\Phi : \mathrm{rank}\ A + \mathrm{rank}\ B = H + r \rbrace $$
(note that $\mathrm{rank}\ A = \mathrm{rank}\ B = H$ is impossible on the fiber: a rank-$H$ factor $B$ is injective, forcing $\mathrm{rank}\ BA = \mathrm{rank}\ A$). The set $\Gamma$ is open in $F_\Phi$, since ranks are lower semicontinuous and their sum is capped; it is dense, since for fixed $B$ the solutions of $BA = \Phi$ form an affine space in which a generic point attains $\mathrm{rank}\ A = H + r - \mathrm{rank}\ B$; the rank pair $(\mathrm{rank}\ A, \mathrm{rank}\ B)$ is locally constant on $\Gamma$, and each piece with a fixed rank pair is a smooth manifold on which the Hessian rank of $L$ is constant and equal to the codimension, so the Li–Wang–Arora hypotheses hold locally there. On the complement $\Sigma = F_\Phi \setminus \Gamma$, where the rank sum drops below $H + r$, a tangent-space count shows the fiber is singular. On the fiber, $\mathrm{tr}\ \nabla^2 L(A, B) = M \Vert B \Vert _F^2 + N \Vert A \Vert _F^2$, whose minimizers over $F_\Phi$ are the weighted-balanced factorizations of rank exactly $r$, those with $N \ A A^\top = M \ B^\top B$ (for $M = N$, the balanced ones $A A^\top = B^\top B$), inside the most degenerate stratum. For $a + b = H + r$, write $\Gamma_{a,b} = \lbrace (A,B) \in F_\Phi : \mathrm{rank}\ A = a,\ \mathrm{rank}\ B = b \rbrace $. Fix $x \in \Gamma_{a,b}$. Let $Z^{a,b,x}$ be the maximal diffusion on this smooth rank-pair component with the Li–Wang–Arora coefficients, started at $x$, and define $\tau_\Sigma = \lim_{j \to \infty} \inf\lbrace t : \mathrm{dist}(Z_t^{a,b,x}, \Sigma) \le j^{-1} \rbrace $ and $\tau_\infty = \lim_{R \to \infty} \inf\lbrace t : \Vert Z_t^{a,b,x} \Vert \ge R \rbrace $. The maximal lifetime is $\tau_\Sigma \wedge \tau_\infty$; this stopped process is the object supplied by the smooth local theory.
- (a) Determine $\mathbb{P}_x(\tau_\Sigma < \tau_\infty)$ for every component and starting point. When this probability is positive, determine the law of the limiting point in $\Sigma$, its rank-pair stratum, and the tail of $\tau_\Sigma$ conditional on $\tau_\Sigma < \tau_\infty$.
- (b) Define the continuation selected by SGD: formulate a martingale problem (or Dirichlet form) on the stratified space $F_\Phi$ whose restriction to every smooth stratum is its Katzenberger diffusion, and prove that rescaled label-noise SGD converges to its unique solution across $\Sigma$. Is the selected process absorbed, reflected, or transmitted at each singular stratum?
- (c) The general form: for real-analytic $L$ whose optimal set is smooth with constant Hessian rank off a codimension-one (in $W_0$) singular locus, construct the stopped diffusion on each smooth component, distinguish singular hitting from escape, and prove a rescaled-SGD limit theorem that selects a well-posed continuation across the singular locus.
There is a sharper point beneath the technicalities. The label-noise drift is built from the Hessian trace, quadratic data, while the learning coefficient $\lambda$ sees all orders. Here $\lambda$ is the central invariant of singular learning theory, Watanabe's real log canonical threshold of the loss at a minimizer: the volume of nearby parameters with loss below $\varepsilon$ scales as $\varepsilon^{\lambda}$ times a power of $\log \tfrac1\varepsilon$, and Bayesian posterior mass concentrates on the parts of the optimal set where $\lambda$ is smallest. The transversal profiles $x^4$ ($\lambda = \tfrac14$) and $x^2 y^2$ ($\lambda = \tfrac12$) both have identically vanishing Hessian at the origin, so the drift ties them while $\lambda$ strictly prefers the quartic. Deciding in an explicit example whether SGD follows the Hessian or the threshold determines whether the $\lambda$-stratification — the partition of the optimal set by the local value of $\lambda$ — is the right stratification for SGD at all. See MAIS-A7, Sections 3.4 and 6.2.
References
- [LWA22] Z. Li, T. Wang, S. Arora, What happens after SGD reaches zero loss? — a mathematical framework, International Conference on Learning Representations, 2022. arXiv:2110.06914
- [K91] G. S. Katzenberger, Solutions of a stochastic differential equation forced onto a manifold by a large drift, Annals of Probability 19 (1991), 1587–1628.
- [DML21] A. Damian, T. Ma, J. D. Lee, Label noise SGD provably prefers flat global minimizers, Advances in Neural Information Processing Systems, 2021. arXiv:2106.06530
Related: MAIS-O77 (the learning-coefficient table for the same fiber) · MAIS-O75 (degenerate-noise selection in the smallest case) · MAIS-O7 (the headline staircase this dynamics should descend).