Leveraging Non-uniformity in First-order Non-convex Optimization
This work addresses convergence bottlenecks in machine learning optimization, offering significant speed-ups for problems like reinforcement learning and supervised learning, though it is incremental in refining existing theoretical frameworks.
The paper tackled the problem of slow convergence in first-order non-convex optimization by introducing non-uniform smoothness and Łojasiewicz inequality, leading to geometry-aware methods that achieve faster convergence rates, such as O(e^{-t}) for policy gradient optimization and linear rates for generalized linear models.
Classical global convergence results for first-order methods rely on uniform smoothness and the Łojasiewicz inequality. Motivated by properties of objective functions that arise in machine learning, we propose a non-uniform refinement of these notions, leading to \emph{Non-uniform Smoothness} (NS) and \emph{Non-uniform Łojasiewicz inequality} (NŁ). The new definitions inspire new geometry-aware first-order methods that are able to converge to global optimality faster than the classical $Ω(1/t^2)$ 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.