Quickly Finding a Benign Region via Heavy Ball Momentum in Non-Convex Optimization

Abstract

The Heavy Ball Method (Polyak, 1964), proposed by Polyak over five decades ago, is a first-order method for optimizing continuous functions. While its stochastic counterpart has proven extremely popular in training deep networks, there are almost no known functions where deterministic Heavy Ball is provably faster than the simple and classical gradient descent algorithm in non-convex optimization. The success of Heavy Ball has thus far eluded theoretical understanding. Our goal is to address this gap, and in the present work we identify two non-convex problems where we provably show that the Heavy Ball momentum helps the iterate to enter a benign region that contains a global optimal point faster. We show that Heavy Ball exhibits simple dynamics that clearly reveal the benefit of using a larger value of momentum parameter for the problems. The first of these optimization problems is the phase retrieval problem, which has useful applications in physical science. The second of these optimization problems is the cubic-regularized minimization, a critical subroutine required by Nesterov-Polyak cubic-regularized method (Nesterov & Polyak (2006)) to find second-order stationary points in general smooth non-convex problems.

Analytical Study of Momentum-Based Acceleration Methods in Paradigmatic High-Dimensional Non-Convex Problems

“…We can now move towards lines (49)(50)(51)(52) and work on simplifying these equa-tions. We start with line (49):…”

“…Where line (56) follows by the definition of response function in y. Moving towards lines (50,51), and carefully taking into account the permutations, we obtain…”

“…Which algorithm is the best in practice seems not to have a simple answer and there are instances where a class of algorithms outperforms the other and vice-versa [26]. Most of the theoretical literature on momentum-based methods concerns convex problems [18,23,24,30,49] and, despite these methods have been successfully applied to a variety of problems, only recently high dimensional non-convex settings have been considered [22,51,52]. Furthermore, with few exceptions [45], the majority of these studies focus on worst-case analysis while empirically one could also be interested in the behaviour of such algorithms on typical instances of the optimization problem, when this is extracted from a probability distribution.…”

“…We apply our equations to the spiked matrix-tensor model which displays a similar phenomenology as the one described in [51] for the phase retrieval problem: all algorithms have two dynamical regimes. First, they navigate in the non-convex landscape and, second, if the signal to noise ratio is strong enough, the dynamics eventually enters in the basin of attraction of the signal and rapidly reaches the bottom of the cost function.…”