Convergence of Markov Chains for Constant Step-size Stochastic Gradient Descent with Separable Functions
David Shirokoff, Philip Zaleski
TL;DR
This work analyzes constant-step SGD for separable non-convex objectives by modeling the iterates as a Markov chain on a general state space. It proves a Doeblin-type decomposition that splits the state space into a uniformly transient region and absorbing rectangles, each carrying a unique invariant measure, and shows that all invariant measures form a convex hull and SGD converges geometrically to a mixture of these invariants. The study demonstrates that diffusion approximations can mispredict long-time behavior, including failing to sample the global minimum and allowing bifurcations between local minima. Together, these results provide a rigorous framework for understanding SGD dynamics beyond diffusion models and suggest directions for extending the theory to broader function classes and dynamical regimes.
Abstract
Stochastic gradient descent (SGD) is a popular algorithm for minimizing objective functions that arise in machine learning. For constant step-sized SGD, the iterates form a Markov chain on a general state space. Focusing on a class of separable (non-convex) objective functions, we establish a "Doeblin-type decomposition," in that the state space decomposes into a uniformly transient set and a disjoint union of absorbing sets. Each of the absorbing sets contains a unique invariant measure, with the set of all invariant measures being the convex hull. Moreover the set of invariant measures are shown to be global attractors to the Markov chain with a geometric convergence rate. The theory is highlighted with examples that show: (1) the failure of the diffusion approximation to characterize the long-time dynamics of SGD; (2) the global minimum of an objective function may lie outside the support of the invariant measures (i.e., even if initialized at the global minimum, SGD iterates will leave); and (3) bifurcations may enable the SGD iterates to transition between two local minima. Key ingredients in the theory involve viewing the SGD dynamics as a monotone iterated function system and establishing a "splitting condition" of Dubins and Freedman 1966 and Bhattacharya and Lee 1988.
