Asymptotic Distribution of Stochastic Mirror Descent Iterates in Average Ensemble Models

The stochastic mirror descent (SMD) algorithm is a general class of training algorithms that utilizes a mirror potential to influence the implicit bias of the training algorithm and includes stochastic gradient descent (SGD) as a special case. In this paper, we explore the performance of the SMD on mean-field ensemble models and generalize earlier results obtained for SGD. The evolution of the distribution of parameters is mapped to a continuous time process in the space of probability distributions. Our main result gives a nonlinear partial differential equation (PDE) to which the continuous time process converges in the asymptotic of large networks. The impact of the mirror potential appears through a multiplicative term that is equal to the inverse of its Hessian and defines a gradient flow over an appropriate Riemannian manifold. We provide numerical simulations which allow us to study and characterize the effect of the mirror potential on the performance of networks trained with SMD for some binary classification problems.