# Leveraging Non-uniformity in First-order Non-convex Optimization

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

## Key facts

- Cited by 44 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 29 references that are in the map.
- First posted on arXiv 13 May 2021.

## Builds on (most cited first, 13 of 13)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Convex Optimization: Algorithms and Complexity](https://mlsanity.com/paper/1405.4980/) (2014)
- [Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/1608.04636/) (2016)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [Why gradient clipping accelerates training: A theoretical justification for adaptivity](https://mlsanity.com/paper/1905.11881/) (2019)
- [On the Global Convergence Rates of Softmax Policy Gradient Methods](https://mlsanity.com/paper/2005.06392/) (2020)
- [Provably Efficient Maximum Entropy Exploration](https://mlsanity.com/paper/1812.02690/) (2018)
- [On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift](https://mlsanity.com/paper/1908.00261/) (2019)
- [Fast Global Convergence of Natural Policy Gradient Methods with Entropy Regularization](https://mlsanity.com/paper/2007.06558/) (2020)
- [Beyond Convexity: Stochastic Quasi-Convex Optimization](https://mlsanity.com/paper/1507.02030/) (2015)
- [Softmax Policy Gradient Methods Can Take Exponential Time to Converge](https://mlsanity.com/paper/2102.11270/) (2021)
- [Accelerating Rescaled Gradient Descent: Fast Optimization of Smooth Functions](https://mlsanity.com/paper/1902.08825/) (2019)

## Built on by (most cited first, 30 of 44)

- [ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models](https://mlsanity.com/paper/2310.10505/) (2023)
- [On the Convergence Rates of Policy Gradient Methods](https://mlsanity.com/paper/2201.07443/) (2022)
- [Policy Mirror Descent for Regularized Reinforcement Learning: A Generalized Framework with Linear Convergence](https://mlsanity.com/paper/2105.11066/) (2021)
- [Stochastic Policy Gradient Methods: Improved Sample Complexity for Fisher-non-degenerate Policies](https://mlsanity.com/paper/2302.01734/) (2023)
- [Linear Convergence of Natural Policy Gradient Methods with Log-Linear Policies](https://mlsanity.com/paper/2210.01400/) (2022)
- [MADE: Exploration via Maximizing Deviation from Explored Regions](https://mlsanity.com/paper/2106.10268/) (2021)
- [A Dual Approach to Constrained Markov Decision Processes with Entropy Regularization](https://mlsanity.com/paper/2110.08923/) (2021)
- [A Novel Framework for Policy Mirror Descent with General Parameterization and Linear Convergence](https://mlsanity.com/paper/2301.13139/) (2023)
- [On the Global Optimum Convergence of Momentum-based Policy Gradient](https://mlsanity.com/paper/2110.10116/) (2021)
- [Beyond Exact Gradients: Convergence of Stochastic Soft-Max Policy Gradient Methods with Entropy Regularization](https://mlsanity.com/paper/2110.10117/) (2021)
- [The Role of Baselines in Policy Gradient Optimization](https://mlsanity.com/paper/2301.06276/) (2023)
- [A Fisher-Rao gradient flow for entropy-regularised Markov decision processes in Polish spaces](https://mlsanity.com/paper/2310.02951/) (2023)
- [Geometry of Polynomial Neural Networks](https://mlsanity.com/paper/2402.00949/) (2024)
- [Stepwise Guided Policy Optimization: Coloring your Incorrect Reasoning in GRPO](https://mlsanity.com/paper/2505.11595/) (2025)
- [Understanding the Effect of Stochasticity in Policy Optimization](https://mlsanity.com/paper/2110.15572/) (2021)
- [Directional Smoothness and Gradient Methods: Convergence and Adaptivity](https://mlsanity.com/paper/2403.04081/) (2024)
- [On the Convergence of Adam under Non-uniform Smoothness: Separability from SGDM and Beyond](https://mlsanity.com/paper/2403.15146/) (2024)
- [QuantFactor REINFORCE: Mining Steady Formulaic Alpha Factors with Variance-bounded REINFORCE](https://mlsanity.com/paper/2409.05144/) (2024)
- [Geometry and convergence of natural policy gradient methods](https://mlsanity.com/paper/2211.02105/) (2022)
- [First-Order Algorithms for Min-Max Optimization in Geodesic Metric Spaces](https://mlsanity.com/paper/2206.02041/) (2022)
- [Stochastic Gradient Succeeds for Bandits](https://mlsanity.com/paper/2402.17235/) (2024)
- [Global Convergence of Natural Policy Gradient with Hessian-aided Momentum Variance Reduction](https://mlsanity.com/paper/2401.01084/) (2024)
- [Towards Statistical and Computational Complexities of Polyak Step Size Gradient Descent](https://mlsanity.com/paper/2110.07810/) (2021)
- [Convergence and Optimality of Policy Gradient Methods in Weakly Smooth Settings](https://mlsanity.com/paper/2111.00185/) (2021)
- [From Gradient Flow on Population Loss to Learning with Stochastic Gradient Descent](https://mlsanity.com/paper/2210.06705/) (2022)
- [Non-Uniform Smoothness for Gradient Descent](https://mlsanity.com/paper/2311.08615/) (2023)
- [Almost sure convergence rates of stochastic gradient methods under gradient domination](https://mlsanity.com/paper/2405.13592/) (2024)
- [Small steps no more: Global convergence of stochastic gradient bandits for arbitrary learning rates](https://mlsanity.com/paper/2502.07141/) (2025)
- [Armijo Line-search Can Make (Stochastic) Gradient Descent Provably Faster](https://mlsanity.com/paper/2503.00229/) (2025)
- [On the Convergence of Policy Mirror Descent with Temporal Difference Evaluation](https://mlsanity.com/paper/2509.18822/) (2025)

## Abstract (arXiv)

Classical global convergence results for first-order methods rely on uniform smoothness and the Lojasiewicz inequality. Motivated by properties of objective functions that arise in machine learning, we propose a non-uniform refinement of these notions, leading to Non-uniform Smoothness (NS) and emphNon-uniform Lojasiewicz inequality (NL). The new definitions inspire new geometry-aware first-order methods that are able to converge to global optimality faster than the classical Omega(1/t2) lower bounds. To illustrate the power of these geometry-aware methods and their corresponding non-uniform analysis, we consider two important problems in machine learning: policy gradient optimization in reinforcement learning (PG), and generalized linear model training in supervised learning (GLM). For PG, we find that normalizing the gradient ascent method can accelerate convergence to O(e-t) while incurring less overhead than existing algorithms. For GLM, we show that geometry-aware normalized gradient descent can also achieve a linear convergence rate, which significantly improves the best known results. We additionally show that the proposed geometry-aware descent methods escape landscape plateaus faster than standard gradient descent. Experimental results are used to illustrate and complement the theoretical findings.