Title:Scaling Limits of Constant-Stepsize SGD at Flat Minima

Abstract:For stochastic gradient descent (SGD) with a constant stepsize $\alpha$, the invariant law of the iterates, centered at a minimizer, describes the behavior of the algorithm over long time horizons. In the strongly convex case, this invariant law has the familiar $\sqrt{\alpha}$ scaling and a Gaussian limit as $\alpha\downarrow 0$. We show that this behavior changes fundamentally for convex objectives $H$ with flat minima and (sub)quadratic tails.
More specifically, we study SGD with Markovian noise generated by a contractive driving chain. For every sufficiently small constant stepsize $\alpha$, we prove existence, uniqueness, and geometric convergence to an augmented invariant law in a Wasserstein distance induced by an $\alpha$-dependent metric. When the minimizer $x_\star$ has local flatness exponent $m\ge2$, meaning that $\nabla^2 H(x)\asymp \lVert x-x_\star\rVert^{m-2} I_d$ as $x\to x_\star$, we obtain a contraction bound with factor $1-c\alpha^{m-1}$, where $c>0$ is a constant. This recovers the factor $1-c\alpha$ in the quadratic case $m=2$. We then analyze the small-stepsize scaling limit. We show that the invariant law concentrates on the scale $\alpha^{1/m}$ and that the rescaled iterates converge weakly to the stationary distribution of the stochastic differential equation $$
dY_t=-h_0(Y_t)\,dt+\Sigma^{1/2}\,dB_t , $$ where $h_0$ is the limiting drift at the minimizer and $\Sigma$ denotes the asymptotic covariance. This recovers the Gaussian limit when $m=2$ and gives generally non-Gaussian stationary limits in the flat case $m>2$. Finally, we give corresponding results for coordinate-separable objectives with unequal flatness exponents.

Bookmark

Bibliographic and Citation Tools

Code, Data and Media Associated with this Article

Demos

Recommenders and Search Tools