Absorbing state dynamics of stochastic gradient descent

Abstract

Stochastic gradient descent (SGD) is a fundamental tool for training deep neural networks across a variety of tasks. In self-supervised learning, different input categories map to distinct manifolds in the embedded neural state space. Accurate classification is achieved by separating these manifolds during learning, akin to a packing problem. We investigate the dynamics of ``neural manifold packing'' by employing a minimal model in which SGD is applied to spherical particles in physical space. In this model, SGD minimizes the system's energy (classification loss) by stochastically reducing overlaps between particles (manifolds). We observe that this process undergoes an absorbing phase transition, prompting us to use the framework of biased random organization (BRO), a nonequilibrium absorbing state model, to describe SGD behavior. We show that BRO dynamics can be approximated by those of particles with linear repulsive interactions under multiplicative anisotropic noise. Thus, for a linear repulsive potential and small kick sizes (learning rates), we find that BRO and SGD become equivalent, converging to the same critical packing fraction $ϕ_c \approx 0.64$, despite the fundamentally different origins of their noise. This equivalence is further supported by the observation that, like BRO, near the critical point, SGD exhibits behavior consistent with the Manna universality class. Above the critical point, SGD exhibits a bias towards flatter minima of the energy landscape, reinforcing the analogy with loss minimization in neural networks.