A Distributed Algorithm for Measure-valued Optimization with Additive Objective

TL;DR

This paper introduces a novel distributed nonparametric algorithm based on a two-layer ADMM for measure-valued optimization problems, applicable in stochastic learning and control contexts like Langevin sampling and neural network training.

Contribution

It extends the Euclidean consensus ADMM to Wasserstein and Sinkhorn consensus ADMM, enabling operator splitting for gradient flows on probability measure manifolds.

Findings

The algorithm effectively solves measure-valued optimization problems.

It generalizes existing ADMM methods to Wasserstein spaces.

Demonstrates applicability in stochastic learning and control scenarios.

Abstract

We propose a distributed nonparametric algorithm for solving measure-valued optimization problems with additive objectives. Such problems arise in several contexts in stochastic learning and control including Langevin sampling from an unnormalized prior, mean field neural network learning and Wasserstein gradient flows. The proposed algorithm comprises a two-layer alternating direction method of multipliers (ADMM). The outer-layer ADMM generalizes the Euclidean consensus ADMM to the Wasserstein consensus ADMM, and to its entropy-regularized version Sinkhorn consensus ADMM. The inner-layer ADMM turns out to be a specific instance of the standard Euclidean ADMM. The overall algorithm realizes operator splitting for gradient flows in the manifold of probability measures.