Accelerated Information Gradient flow

TL;DR

This paper introduces a unified accelerated gradient flow framework in probability space for efficient Bayesian sampling, demonstrating convergence and proposing practical algorithms with superior performance in Bayesian inference tasks.

Contribution

It develops a novel accelerated gradient flow framework for multiple information metrics and introduces efficient discrete algorithms with a restart technique for Bayesian inverse problems.

Findings

Proves convergence for Fisher-Rao and Wasserstein-2 gradient flows.

Proposes a sampling-efficient discrete-time algorithm with restart.

Demonstrates improved performance in Bayesian logistic regression and neural networks.

Abstract

We present a framework for Nesterov's accelerated gradient flows in probability space to design efficient mean-field Markov chain Monte Carlo (MCMC) algorithms for Bayesian inverse problems. Here four examples of information metrics are considered, including Fisher-Rao metric, Wasserstein-2 metric, Kalman-Wasserstein metric and Stein metric. For both Fisher-Rao and Wasserstein-2 metrics, we prove convergence properties of accelerated gradient flows. In implementations, we propose a sampling-efficient discrete-time algorithm for Wasserstein-2, Kalman-Wasserstein and Stein accelerated gradient flows with a restart technique. We also formulate a kernel bandwidth selection method, which learns the gradient of logarithm of density from Brownian-motion samples. Numerical experiments, including Bayesian logistic regression and Bayesian neural network, show the strength of the proposed methods compared with state-of-the-art algorithms.