Riemannian Momentum Tracking: Distributed Optimization with Momentum on Compact Submanifolds

Abstract

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.