Particle-based Variational Inference with Generalized Wasserstein Gradient Flow

TL;DR

This paper introduces GWG, a particle-based variational inference framework using generalized Wasserstein gradient flow, offering improved flexibility, convergence guarantees, and adaptive acceleration for efficient inference.

Contribution

The paper proposes a novel GWG framework that generalizes Wasserstein gradient flow with convex regularizers, enhancing flexibility and convergence in particle-based variational inference.

Findings

GWG demonstrates strong convergence guarantees.

Adaptive GWG accelerates convergence effectively.

Proven effectiveness on simulated and real data.

Abstract

Particle-based variational inference methods (ParVIs) such as Stein variational gradient descent (SVGD) update the particles based on the kernelized Wasserstein gradient flow for the Kullback-Leibler (KL) divergence. However, the design of kernels is often non-trivial and can be restrictive for the flexibility of the method. Recent works show that functional gradient flow approximations with quadratic form regularization terms can improve performance. In this paper, we propose a ParVI framework, called generalized Wasserstein gradient descent (GWG), based on a generalized Wasserstein gradient flow of the KL divergence, which can be viewed as a functional gradient method with a broader class of regularizers induced by convex functions. We show that GWG exhibits strong convergence guarantees. We also provide an adaptive version that automatically chooses Wasserstein metric to accelerate convergence. In experiments, we demonstrate the effectiveness and efficiency of the proposed framework on both simulated and real data problems.