Scaling Limits of Constant-Stepsize SGD at Flat Minima
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 \geq 2\), 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 \[\mathrm{d}Y_t=-h_0(Y_t)\,\mathrm{d}t+\Sigma^{1/2}\,\mathrm{d}B_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.