Wasserstein Consensus ADMM

TL;DR

This paper introduces Wasserstein consensus ADMM and Sinkhorn consensus ADMM algorithms for measure-valued convex optimization, enabling distributed solutions for Wasserstein gradient flows with applications in stochastic prediction and learning.

Contribution

It generalizes Euclidean ADMM to the space of probability measures with a two-layer algorithm structure, including an entropic regularization variant for improved computation.

Findings

Effective in solving Wasserstein gradient flows

Suitable for distributed computation in measure spaces

Demonstrated with numerical examples

Abstract

We introduce Wasserstein consensus alternating direction method of multipliers (ADMM) and its entropic-regularized version: Sinkhorn consensus ADMM, to solve measure-valued optimization problems with convex additive objectives. Several problems of interest in stochastic prediction and learning can be cast in this form of measure-valued convex additive optimization. The proposed algorithm generalizes a variant of the standard Euclidean ADMM to the space of probability measures but departs significantly from its Euclidean counterpart. In particular, we derive a two layer ADMM algorithm wherein the outer layer is a variant of consensus ADMM on the space of probability measures while the inner layer is a variant of Euclidean ADMM. The resulting computational framework is particularly suitable for solving Wasserstein gradient flows via distributed computation. We demonstrate the proposed framework using illustrative numerical examples.