# Nonconvex Stochastic Bregman Proximal Gradient Method with Application to Deep Learning

arXiv: https://arxiv.org/abs/2306.14522 · HTML: https://mlsanity.com/paper/2306.14522/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 95 references that are in the map.
- First posted on arXiv 26 Jun 2023.

## Builds on (most cited first, 30 of 45)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [A ConvNet for the 2020s](https://mlsanity.com/paper/2201.03545/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Convex Optimization: Algorithms and Complexity](https://mlsanity.com/paper/1405.4980/) (2014)
- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [Lower Bounds for Non-Convex Stochastic Optimization](https://mlsanity.com/paper/1912.02365/) (2019)
- [AdaBelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients](https://mlsanity.com/paper/2010.07468/) (2020)
- [Adaptive Gradient Methods with Dynamic Bound of Learning Rate](https://mlsanity.com/paper/1902.09843/) (2019)
- [Nonconvex Optimization Meets Low-Rank Matrix Factorization: An Overview](https://mlsanity.com/paper/1809.09573/) (2018)
- [Mini-batch Stochastic Approximation Methods for Nonconvex Stochastic Composite Optimization](https://mlsanity.com/paper/1308.6594/) (2013)
- [Katyusha: The First Direct Acceleration of Stochastic Gradient Methods](https://mlsanity.com/paper/1603.05953/) (2016)
- [Towards Theoretically Understanding Why SGD Generalizes Better Than ADAM in Deep Learning](https://mlsanity.com/paper/2010.05627/) (2020)
- [Matrix Completion has No Spurious Local Minimum](https://mlsanity.com/paper/1605.07272/) (2016)
- [A Closer Look at Deep Learning Heuristics: Learning rate restarts, Warmup and Distillation](https://mlsanity.com/paper/1810.13243/) (2018)
- [Relatively-Smooth Convex Optimization by First-Order Methods, and Applications](https://mlsanity.com/paper/1610.05708/) (2016)
- [Stochastic Compositional Gradient Descent: Algorithms for Minimizing Compositions of Expected-Value Functions](https://mlsanity.com/paper/1411.3803/) (2014)
- [A Geometric Analysis of Phase Retrieval](https://mlsanity.com/paper/1602.06664/) (2016)
- [Closing the Generalization Gap of Adaptive Gradient Methods in Training Deep Neural Networks](https://mlsanity.com/paper/1806.06763/) (2018)
- [On the Ineffectiveness of Variance Reduced Optimization for Deep Learning](https://mlsanity.com/paper/1812.04529/) (2018)
- [First Order Methods beyond Convexity and Lipschitz Gradient Continuity with Applications to Quadratic Inverse Problems](https://mlsanity.com/paper/1706.06461/) (2017)

## Built on by (most cited first, 7 of 7)

- [Taming Nonconvex Stochastic Mirror Descent with General Bregman Divergence](https://mlsanity.com/paper/2402.17722/) (2024)
- [Adam-family Methods with Decoupled Weight Decay in Deep Learning](https://mlsanity.com/paper/2310.08858/) (2023)
- [On exploration of an interior mirror descent flow for stochastic nonconvex constrained problem](https://mlsanity.com/paper/2507.15264/) (2025)
- [Shuffling the Stochastic Mirror Descent via Dual Lipschitz Continuity and Kernel Conditioning](https://mlsanity.com/paper/2603.16042/) (2026)
- [Non-KKT Accumulation in Entropic Mirror Descent](https://mlsanity.com/paper/2608.01658/) (2026)
- [On the Tension Between Optimality and Adversarial Robustness in Policy Optimization](https://mlsanity.com/paper/2512.01228/) (2025)
- [Establishing Boundary KKT Convergence of Mirror Descent through Reparameterization](https://mlsanity.com/paper/2608.07248/) (2026)

## Abstract (arXiv)

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 like quadratic inverse problems and neural network training, 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 MSBPG to the training of deep neural networks, utilizing a polynomial kernel function to ensure smooth adaptivity of the loss function. Experimental results on benchmark datasets confirm the effectiveness and robustness of MSBPG in training neural networks. Given its negligible additional computational cost compared to SGD in large-scale optimization, MSBPG shows promise as a universal open-source optimizer for future applications.