Variational Transport: A Convergent Particle-BasedAlgorithm for Distributional Optimization
This work is significant for researchers and practitioners working with distributional optimization problems, offering a new algorithm with theoretical convergence guarantees. It provides an incremental improvement in the methodology for these problems.
This paper addresses the problem of minimizing a functional defined over a family of probability distributions, a common task in machine learning and statistics. The authors propose Variational Transport, a particle-based algorithm that approximates Wasserstein gradient descent. They prove that under functional Polyak-Łojasiewicz and smoothness conditions, their algorithm converges linearly to the global minimum, with statistical error decaying sublinearly as particle count increases.
We consider the optimization problem of minimizing a functional defined over a family of probability distributions, where the objective functional is assumed to possess a variational form. Such a distributional optimization problem arises widely in machine learning and statistics, with Monte-Carlo sampling, variational inference, policy optimization, and generative adversarial network as examples. For this problem, we propose a novel particle-based algorithm, dubbed as variational transport, which approximately performs Wasserstein gradient descent over the manifold of probability distributions via iteratively pushing a set of particles. Specifically, we prove that moving along the geodesic in the direction of functional gradient with respect to the second-order Wasserstein distance is equivalent to applying a pushforward mapping to a probability distribution, which can be approximated accurately by pushing a set of particles. Specifically, in each iteration of variational transport, we first solve the variational problem associated with the objective functional using the particles, whose solution yields the Wasserstein gradient direction. Then we update the current distribution by pushing each particle along the direction specified by such a solution. By characterizing both the statistical error incurred in estimating the Wasserstein gradient and the progress of the optimization algorithm, we prove that when the objective function satisfies a functional version of the Polyak-Łojasiewicz (PL) (Polyak, 1963) and smoothness conditions, variational transport converges linearly to the global minimum of the objective functional up to a certain statistical error, which decays to zero sublinearly as the number of particles goes to infinity.