Riemannian Momentum Tracking: Distributed Optimization With Momentum on Compact Submanifolds

Gradient descent with momentum has been widely applied in various signal processing and machine learning tasks, demonstrating a notable empirical advantage over standard gradient descent. However, momentum-based distributed Riemannian algorithms have been only scarcely explored. In this paper, we propose Riemannian Momentum Tracking (RMTracking), a decentralized optimization algorithm with momentum over a compact submanifold. Given the non-convex nature of compact submanifolds, the objective function, composed of a finite sum of smooth (possibly non-convex) local functions, is minimized across agents in an undirected and connected network graph. With a constant step-size, we establish an$\mathcal {O}(\frac{1-\beta }{K})$convergence rate of the Riemannian gradient average for any momentum weight$\beta \in [0,1)$. Especially, RMTracking can achieve a convergence rate of$\mathcal {O}(\frac{1-\beta }{K})$to a stationary point when the step-size is sufficiently small. To best of our knowledge, RMTracking is the first decentralized algorithm to achieve exact convergence that is$\frac{1}{1-\beta }$times faster than other related algorithms. Finally, we verify these theoretical claims through numerical experiments on eigenvalue problems.