Polyak's Heavy Ball Method Achieves Accelerated Local Rate of Convergence under Polyak-Lojasiewicz Inequality (2410.16849v1)

Abstract: In this work, we consider the convergence of Polyak's heavy ball method, both in continuous and discrete time, on a non-convex objective function. We recover the convergence rates derived in [Polyak, U.S.S.R. Comput. Math. and Math. Phys., 1964] for strongly convex objective functions, assuming only validity of the Polyak-Lojasiewicz inequality. In continuous time our result holds for all initializations, whereas in the discrete time setting we conduct a local analysis around the global minima. Our results demonstrate that the heavy ball method does, in fact, accelerate on the class of objective functions satisfying the Polyak-Lojasiewicz inequality. This holds even in the discrete time setting, provided the method reaches a neighborhood of the global minima. Instead of the usually employed Lyapunov-type arguments, our approach leverages a new differential geometric perspective of the Polyak-Lojasiewicz inequality proposed in [Rebjock and Boumal, Math. Program., 2024].

• The paper proves that Polyak's heavy ball method attains an optimal 2√μ convergence rate under the PL inequality in both continuous and discrete settings.
• It employs a differential geometric approach by decomposing the dynamics into tangential and normal components, highlighting the strong attraction in the normal direction.
• Numerical experiments validate the theoretical findings, demonstrating the method’s efficacy for efficient optimization in machine learning.

Overview of Polyak's Heavy Ball Method and PL-Inequality

The paper "Polyak's Heavy Ball Method Achieves Accelerated Local Rate of Convergence under Polyak-Łojasiewicz Inequality" by Sebastian Kassing and Simon Weissmann explores the convergence properties of the heavy ball method (HBM) in the context of non-convex optimization. It specifically investigates the applicability of the Polyak-Łojasiewicz (PL) inequality as a relaxation of the strong convexity assumption commonly employed in optimization theory.

Background and Motivation

In optimization, achieving faster convergence rates is crucial, especially for high-dimensional non-convex problems encountered in machine learning. Traditional methods like gradient descent are often enhanced by momentum-based techniques, with Polyak's heavy ball method being a prominent example. These methods introduce an additional term that accelerates convergence by utilizing past gradient information.

The PL condition, initially proposed by Polyak and Łojasiewicz, serves as a less restrictive alternative to strong convexity, making it particularly suitable for non-convex loss functions typical in machine learning applications. The authors aim to establish that HBM achieves accelerated convergence under the PL condition.

Main Contributions

1. Convergence Analysis in Continuous Time: The paper revisits Polyak's original convergence rates for HBM in continuous time. Under the PL inequality, the authors prove that the convergence rate for the heavy ball ordinary differential equation (ODE) can achieve the optimal $2\sqrt{\mu}$, where $\mu$ is the PL constant. This result is significant as it indicates that the PL inequality suffices to obtain the accelerated rates previously attributed only to strongly convex functions.
2. Discrete Time Convergence: For the discrete version of HBM, the authors derive local convergence rates. They demonstrate that if the iterates reach a vicinity of the global minimizer, the convergence rate is consistent with the accelerated rate observed in the strongly convex case, contingent on the PL condition being fulfilled.
3. Geometric Perspective: A novel perspective is introduced by leveraging the geometry induced by the PL inequality. This approach allows the decomposition of the optimization problem into tangential and normal components relative to the manifold of global minima. The analysis shows that the normal direction is strongly attracting, thus facilitating accelerated convergence.
4. Numerical Results and Discussion: The authors validate their theoretical findings with numerical experiments that corroborate the improved convergence rates under the PL condition.

Implications and Future Directions

This paper significantly broadens the applicability of fast convergence rates to non-convex settings. The implications are noteworthy for machine learning models, where non-convexity is standard. The heavy ball method, when combined with the PL inequality, offers an effective mechanism for optimizing such models with favorable computational efficiency.

Future directions could explore heterogeneous function classes where local minima are polyhedral or have more diverse geometrical structures. Extending these results to stochastic settings or other momentum-based methods, such as Nesterov's accelerated gradient, might also yield fruitful insights. Furthermore, understanding the interplay between the geometry imposed by the PL condition and various optimization landscapes could guide the design of new optimization algorithms.

Conclusion

This paper successfully establishes the accelerated convergence of Polyak's heavy ball method under the PL inequality, highlighting its viability beyond strongly convex functions. By providing a differential geometric viewpoint, the authors enhance our understanding of how these mathematical properties can be leveraged to achieve optimal convergence rates in complex, non-convex scenarios. This work is a pivotal advancement for researchers and practitioners involved in the optimization of machine learning models.