The stochastic heavy ball method (SHB), also known as stochastic gradient descent (SGD) with Polyak's momentum, is widely used in training neural networks. However, despite the remarkable success of such algorithm in practice, its theoretical characterization remains limited. In this paper, we focus on neural networks with two and three layers and provide a rigorous understanding of the properties of the solutions found by SHB: \emph{(i)} stability after dropping out part of the neurons, \emph{(ii)} connectivity along a low-loss path, and \emph{(iii)} convergence to the global optimum. To achieve this goal, we take a mean-field view and relate the SHB dynamics to a certain partial differential equation in the limit of large network widths. This mean-field perspective has inspired a recent line of work focusing on SGD while, in contrast, our paper considers an algorithm with momentum. More specifically, after proving existence and uniqueness of the limit differential equations, we show convergence to the global optimum and give a quantitative bound between the mean-field limit and the SHB dynamics of a finite-width network. Armed with this last bound, we are able to establish the dropout-stability and connectivity of SHB solutions.

翻译：随机重球方法（SHB），也称为带Polyak动量的随机梯度下降（SGD），广泛应用于神经网络训练。然而，尽管该算法在实践中取得了显著成功，其理论刻画仍相对有限。本文聚焦于两层及三层神经网络，对SHB解的以下性质进行了严格解析：（i）部分神经元丢弃后的稳定性，（ii）沿低损失路径的连通性，以及（iii）全局最优解的收敛性。为实现这一目标，我们采用平均场视角，将SHB动力学与网络宽度趋于无穷时的特定偏微分方程建立关联。该平均场方法近期催生了聚焦于SGD的研究路线，而本文则考虑带动量算法。具体而言，在证明极限微分方程的存在唯一性后，我们展示了全局最优收敛性，并给出了有限宽度网络的平均场极限与SHB动力学之间的定量界。基于该定量界，我们得以建立SHB解的失稳性与连通性。