Kernel Stein Discrepancy Descent

TL;DR

This paper introduces KSD Descent, a new deterministic, score-based sampling method leveraging Wasserstein gradient flow of Kernel Stein Discrepancy, offering a parameter-free optimization approach for approximating target distributions.

Contribution

It proposes KSD Descent, a novel particle-based sampling algorithm using Wasserstein gradient flow and robust optimization, with theoretical analysis and practical demonstrations.

Findings

KSD Descent effectively approximates target distributions.

It enables parameter-free optimization with L-BFGS.

The method can get stuck in local minima in some cases.

Abstract

Among dissimilarities between probability distributions, the Kernel Stein Discrepancy (KSD) has received much interest recently. We investigate the properties of its Wasserstein gradient flow to approximate a target probability distribution $\pi$ on $\mathbb{R}^d$, known up to a normalization constant. This leads to a straightforwardly implementable, deterministic score-based method to sample from $\pi$, named KSD Descent, which uses a set of particles to approximate $\pi$. Remarkably, owing to a tractable loss function, KSD Descent can leverage robust parameter-free optimization schemes such as L-BFGS; this contrasts with other popular particle-based schemes such as the Stein Variational Gradient Descent algorithm. We study the convergence properties of KSD Descent and demonstrate its practical relevance. However, we also highlight failure cases by showing that the algorithm can get stuck in spurious local minima.