Continuous-time Riemannian SGD and SVRG Flows on Wasserstein Probabilistic Space

2024-01-24

arXiv.org

AI Summary

The absence of rigorous continuous-time stochastic optimization dynamics on the Wasserstein space has hindered theoretical analysis and algorithm design for probability measure optimization. Method: We generalize classical Euclidean stochastic optimization to the Riemannian manifold of probability measures endowed with Wasserstein geometry. Specifically, we formulate Riemannian stochastic gradient descent (Riemannian SGD) and Riemannian stochastic variance-reduced gradient (Riemannian SVRG) as stochastic differential equations (SDEs), characterize their measure-valued trajectories via the Fokker–Planck equation, and analyze convergence through KL-divergence minimization and Langevin dynamics. Contribution/Results: We establish, for the first time, well-posed continuous-time flows for both Riemannian SGD and Riemannian SVRG on the Wasserstein space. We prove that their convergence rates match the optimal theoretical bounds of their Euclidean counterparts—namely, $O(1/t)$ for SGD and linear convergence for SVRG. This work provides a theoretically grounded, continuous-time optimization framework applicable to Bayesian sampling, variational inference, and related tasks.

Technology Category

• Reasoning under Uncertainty: Stochastic Optimization
• Machine Learning: Learning with Manifolds
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web
• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods

Abstract

Recently, optimization on the Riemannian manifold has provided new insights to the optimization community. In this regard, the manifold taken as the probability measure metric space equipped with the second-order Wasserstein distance is of particular interest, since optimization on it can be linked to practical sampling processes. In general, the standard (continuous) optimization method on Wasserstein space is Riemannian gradient flow (i.e., Langevin dynamics when minimizing KL divergence). In this paper, we aim to enrich the continuous optimization methods in the Wasserstein space, by extending the gradient flow on it into the stochastic gradient descent (SGD) flow and stochastic variance reduction gradient (SVRG) flow. The two flows in Euclidean space are standard continuous stochastic methods, while their Riemannian counterparts are unexplored. By leveraging the property of Wasserstein space, we construct stochastic differential equations (SDEs) to approximate the corresponding discrete dynamics of desired Riemannian stochastic methods in Euclidean space. Then, our probability measures flows are obtained by the Fokker-Planck equation. Finally, the convergence rates of our Riemannian stochastic flows are proven, which match the results in Euclidean space.

Problem

• Extends Riemannian SGD and SVRG flows to Wasserstein probabilistic space
• Constructs SDEs to approximate discrete Euclidean stochastic optimization dynamics
• Establishes convergence rates matching Euclidean settings for proposed flows

Innovation

• Extends SGD and SVRG flows to Wasserstein space
• Constructs SDEs to approximate discrete Riemannian dynamics
• Derives flows via Fokker-Planck equation with convergence guarantees