# Taming Nonconvex Stochastic Mirror Descent with General Bregman Divergence arXiv: https://arxiv.org/abs/2402.17722 · HTML: https://mlsanity.com/paper/2402.17722/ ## Key facts - Cited by 11 papers in the map (data through 14 Oct 2026). - Builds on 58 of its 94 references that are in the map. - First posted on arXiv 27 Feb 2024. ## Builds on (most cited first, 30 of 58) - [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017) - [Deep Learning with Differential Privacy](https://mlsanity.com/paper/1607.00133/) (2016) - [Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming](https://mlsanity.com/paper/1309.5549/) (2013) - [Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/1608.04636/) (2016) - [Why gradient clipping accelerates training: A theoretical justification for adaptivity](https://mlsanity.com/paper/1905.11881/) (2019) - [Lower Bounds for Non-Convex Stochastic Optimization](https://mlsanity.com/paper/1912.02365/) (2019) - [Global Convergence of Policy Gradient Methods for the Linear Quadratic Regulator](https://mlsanity.com/paper/1801.05039/) (2018) - [Deep Learning without Poor Local Minima](https://mlsanity.com/paper/1605.07110/) (2016) - [Stochastic model-based minimization of weakly convex functions](https://mlsanity.com/paper/1803.06523/) (2018) - [Policy Mirror Descent for Reinforcement Learning: Linear Convergence, New Sampling Complexity, and Generalized Problem Classes](https://mlsanity.com/paper/2102.00135/) (2021) - [Mini-batch Stochastic Approximation Methods for Nonconvex Stochastic Composite Optimization](https://mlsanity.com/paper/1308.6594/) (2013) - [Differentially Private Empirical Risk Minimization Revisited: Faster and More General](https://mlsanity.com/paper/1802.05251/) (2018) - [On the Convergence Rates of Policy Gradient Methods](https://mlsanity.com/paper/2201.07443/) (2022) - [Variational Policy Gradient Method for Reinforcement Learning with General Utilities](https://mlsanity.com/paper/2007.02151/) (2020) - [Relatively-Smooth Convex Optimization by First-Order Methods, and Applications](https://mlsanity.com/paper/1610.05708/) (2016) - [Independent Policy Gradient Methods for Competitive Reinforcement Learning](https://mlsanity.com/paper/2101.04233/) (2021) - [Policy Mirror Descent for Regularized Reinforcement Learning: A Generalized Framework with Linear Convergence](https://mlsanity.com/paper/2105.11066/) (2021) - [Private Stochastic Convex Optimization: Optimal Rates in ℓ1 Geometry](https://mlsanity.com/paper/2103.01516/) (2021) - [Revisiting Gradient Clipping: Stochastic bias and tight convergence guarantees](https://mlsanity.com/paper/2305.01588/) (2023) - [Unified Optimal Analysis of the (Stochastic) Gradient Method](https://mlsanity.com/paper/1907.04232/) (2019) - [On the Convergence and Sample Efficiency of Variance-Reduced Policy Gradient Method](https://mlsanity.com/paper/2102.08607/) (2021) - [Boltzmann Exploration Done Right](https://mlsanity.com/paper/1705.10257/) (2017) - [Efficient Private ERM for Smooth Objectives](https://mlsanity.com/paper/1703.09947/) (2017) - [Efficiency of minimizing compositions of convex functions and smooth maps](https://mlsanity.com/paper/1605.00125/) (2016) - [Linear Coupling: An Ultimate Unification of Gradient and Mirror Descent](https://mlsanity.com/paper/1407.1537/) (2014) - [Stochastic Policy Gradient Methods: Improved Sample Complexity for Fisher-non-degenerate Policies](https://mlsanity.com/paper/2302.01734/) (2023) - [Non-Euclidean Differentially Private Stochastic Convex Optimization: Optimal Rates in Linear Time](https://mlsanity.com/paper/2103.01278/) (2021) - [High Probability Convergence of Stochastic Gradient Methods](https://mlsanity.com/paper/2302.14843/) (2023) - [On the Linear convergence of Natural Policy Gradient Algorithm](https://mlsanity.com/paper/2105.01424/) (2021) - [Differentially Private Stochastic Optimization: New Results in Convex and Non-Convex Settings](https://mlsanity.com/paper/2107.05585/) (2021) ## Built on by (most cited first, 11 of 11) - [From Gradient Clipping to Normalization for Heavy Tailed SGD](https://mlsanity.com/paper/2410.13849/) (2024) - [Can SGD Handle Heavy-Tailed Noise?](https://mlsanity.com/paper/2508.04860/) (2025) - [Understanding Stochastic Natural Gradient Variational Inference](https://mlsanity.com/paper/2406.01870/) (2024) - [Natural Gradient VI: Guarantees for Non-Conjugate Models](https://mlsanity.com/paper/2510.19163/) (2025) - [Sparse Training of Neural Networks based on Multilevel Mirror Descent](https://mlsanity.com/paper/2602.03535/) (2026) - [Convergence of Policy Mirror Descent Beyond Compatible Function Approximation](https://mlsanity.com/paper/2502.11033/) (2025) - [Leveraging Machine Unlearning for Cost-Efficient Preference Alignment](https://mlsanity.com/paper/2504.06659/) (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) - [A New Kernel Regularity Condition for Distributed Mirror Descent: Broader Coverage and Simpler Analysis](https://mlsanity.com/paper/2603.12838/) (2026) - [Convergence Guarantees of Gradient Descent for Neural Networks via Generalized Lipschitz Smoothness](https://mlsanity.com/paper/2608.11479/) (2026) ## Abstract (arXiv) This paper revisits the convergence of Stochastic Mirror Descent (SMD) in the contemporary nonconvex optimization setting. Existing results for batch-free nonconvex SMD restrict the choice of the distance generating function (DGF) to be differentiable with Lipschitz continuous gradients, thereby excluding important setups such as Shannon entropy. In this work, we present a new convergence analysis of nonconvex SMD supporting general DGF, that overcomes the above limitations and relies solely on the standard assumptions. Moreover, our convergence is established with respect to the Bregman Forward-Backward envelope, which is a stronger measure than the commonly used squared norm of gradient mapping. We further extend our results to guarantee high probability convergence under sub-Gaussian noise and global convergence under the generalized Bregman Proximal Polyak-Lojasiewicz condition. Additionally, we illustrate the advantages of our improved SMD theory in various nonconvex machine learning tasks by harnessing nonsmooth DGFs. Notably, in the context of nonconvex differentially private (DP) learning, our theory yields a simple algorithm with a (nearly) dimension-independent utility bound. For the problem of training linear neural networks, we develop provably convergent stochastic algorithms.