Nonconvex Stochastic Bregman Proximal Gradient Method for Nonconvex Composite Problems (2306.14522v4)

Abstract: Stochastic gradient methods for minimizing nonconvex composite objective functions typically rely on the Lipschitz smoothness of the differentiable part, but this assumption fails in many important problem classes, leading to instability of the algorithms in both theory and practice. To address this, we propose a family of stochastic Bregman proximal gradient (SBPG) methods that only require smooth adaptivity. SBPG replaces the quadratic approximation in SGD with a Bregman proximity measure, offering a better approximation model that handles non-Lipschitz gradients in nonconvex objectives. We establish the convergence properties of vanilla SBPG and show it achieves optimal sample complexity in the nonconvex setting. Experimental results on quadratic inverse problems demonstrate SBPG's robustness in terms of stepsize selection and sensitivity to the initial point. Furthermore, we introduce a momentum-based variant, MSBPG, which enhances convergence by relaxing the mini-batch size requirement while preserving the optimal oracle complexity. We apply a polynomial kernel function based MBPG to the loss function with polynomial growth. Experimental results on benchmark datasets confirm the effectiveness and robustness of MSBPG. Given its negligible additional computational cost compared to SGD in large-scale optimization, MSBPG shows promise as a universal optimizer for future applications.