[Paper Review] A Unified Particle-Optimization Framework for Scalable Bayesian Sampling

Changyou Chen, Ruiyi Zhang+3 more|arXiv (Cornell University)|May 29, 2018

Markov Chains and Monte Carlo Methods33 references40 citations

TL;DR

The paper presents a unified particle-optimization framework on Wasserstein gradient flows that unifies SG-MCMC and SVGD, enabling efficient, scalable Bayesian sampling via particle methods.

ABSTRACT

There has been recent interest in developing scalable Bayesian sampling methods such as stochastic gradient MCMC (SG-MCMC) and Stein variational gradient descent (SVGD) for big-data analysis. A standard SG-MCMC algorithm simulates samples from a discrete-time Markov chain to approximate a target distribution, thus samples could be highly correlated, an undesired property for SG-MCMC. In contrary, SVGD directly optimizes a set of particles to approximate a target distribution, and thus is able to obtain good approximations with relatively much fewer samples. In this paper, we propose a principle particle-optimization framework based on Wasserstein gradient flows to unify SG-MCMC and SVGD, and to allow new algorithms to be developed. Our framework interprets SG-MCMC as particle optimization on the space of probability measures, revealing a strong connection between SG-MCMC and SVGD. The key component of our framework is several particle-approximate techniques to efficiently solve the original partial differential equations on the space of probability measures. Extensive experiments on both synthetic data and deep neural networks demonstrate the effectiveness and efficiency of our framework for scalable Bayesian sampling.

Motivation & Objective

• Motivate scalable Bayesian sampling for big data using SG-MCMC and SVGD.
• Propose a particle-optimization framework on Wasserstein gradient flows to unify SG-MCMC and SVGD.
• Develop practical particle-approximation techniques (discrete gradient flows and blob methods) to solve the continuous formulations.
• Demonstrate effectiveness and efficiency through synthetic data and deep neural network experiments.

Proposed method

• Formulate SG-MCMC and SVGD as Wasserstein gradient flows on the space of probability measures.
• Introduce particle-approximation techniques: discrete gradient flows and blob methods, to solve the resulting optimization problems.
• Derive updates for SVGD as a blob-method-based particle evolution that recovers the standard SVGD updates.
• Derive SG-MCMC particle updates by approximating the Jordan–Kinderlehrer–Otto scheme and tackling the W2 term with entropy-regularized transport.
• Provide practical gradient-based updates for particles using approximations of the energy function and transport costs.

Experimental results

Research questions

• RQ1 Can SG-MCMC and SVGD be viewed as special cases of a unified particle-optimization framework on Wasserstein gradient flows?
• RQ2 What particle-approximation techniques enable scalable, efficient solutions for SG-MCMC and SVGD in high dimensions?
• RQ3 How can discrete gradient flows and blob methods be applied to derive practical particle updates for scalable Bayesian sampling?

Key findings

• A unified particle-optimization framework is able to encompass SG-MCMC and SVGD as special cases on Wasserstein gradient flows.
• Discrete gradient-flow and blob-method approaches provide scalable particle-based solutions to the FP equations governing the samplers.
• SVGD can be interpreted as a WGF with a specific velocity field that leads to its standard particle updates.
• SG-MCMC can be implemented via a discrete-gradient-flow scheme leading to practical particle updates aligned with the JKO scheme.
• Experiments on synthetic data and deep neural networks demonstrate the proposed framework’s effectiveness and efficiency for scalable Bayesian sampling.