Title: High Probability Bounds for Non-Convex Stochastic Optimization with Momentum
Primary area: learning theory
Keywords: Momentum, nonconvex learning, generalization

Abstract: Stochastic gradient descent with momentum (SGDM) has been widely used in machine learning. However, in non-convex domains, high probability learning bounds for SGDM are scarce. In this paper, we provide high probability convergence bounds and generalization bounds for SGDM. Firstly, we establish these bounds for the gradient norm in the general non-convex case. The derived convergence bounds are tighter than the theoretical results of related work, and to our best knowledge, the derived generalization bounds are the first ones for SGDM. Then, if the Polyak-{\L}ojasiewicz condition is satisfied, we establish these bounds for the error of the function value, instead of the gradient norm. Moreover, the derived learning bounds have faster rates than the general non-convex case. Finally,  we further provide sharper generalization bounds by considering a mild Bernstein condition on the gradient. In the case of low noise, their learning rates can reach $\widetilde{\mathcal{O}}(1/n^2)$, where $n$ is the sample size. Overall, we relatively systematically investigate the high probability learning bounds for non-convex SGDM.