Beyond the Hessian Edge: The Stochastic Stability Cocycle of Mini-Batch SGD
Abstract
We study the training-time stochastic stability cocycle of multi-pass mini-batch \sgd\ in the Gaussian-design / affine-row-regularizer specialization of fixed-rank proportional-asymptotic multi-index models. Our main result is a finite-horizon replicated tangent \dmft: by transporting finitely many common-noise tangent replicas through the same batch sequence, we obtain a closed deterministic kernel law for replica--replica, replica--iterate, and replica--teacher overlaps. This yields an exact deterministic law for the characteristic (q)-volume growth, defined by the log-determinant growth of the tangent Gram matrix and equivalently interpretable as exterior-power growth or accumulated QR growth of the Jacobian cocycle. Under an asymptotically stationary regime, the finite-time law converges to a top stationary exponent whose zero-crossing defines the stochastic edge. We also compare the Poissonian mini-batch theory with its Brownian \sgf\ reference limit and derive a finite-horizon jump-cumulant correction for the characteristic (q)-volume law, which vanishes exactly in the affine-linear learner-side regime. In the rank-one logistic model, the tangent dynamics split into an explicit signal channel and an orthogonal bulk channel, yielding a closed one-dimensional critical equation for the stochastic edge.