Sampling in Unit Time with Kernel Fisher-Rao Flow
Abstract
We introduce a new mean-field ODE and corresponding interacting particle systems (IPS) for sampling from an unnormalized target density. The IPS are gradient-free, available in closed form, and only require the ability to sample from a reference density and compute the (unnormalized) target-to-reference density ratio. The mean-field ODE is obtained by solving a Poisson equation for a velocity field that transports samples along the geometric mixture of the two densities, which is the path of a particular Fisher-Rao gradient flow. We employ a RKHS ansatz for the velocity field, which makes the Poisson equation tractable and enables discretization of the resulting mean-field ODE over finite samples. The mean-field ODE can be additionally be derived from a discrete-time perspective as the limit of successive linearizations of the Monge-Amp\`ere equations within a framework known as sample-driven optimal transport. We introduce a stochastic variant of our approach and demonstrate empirically that our IPS can produce high-quality samples from varied target distributions, outperforming comparable gradient-free particle systems and competitive with gradient-based alternatives.
Cite
@article{arxiv.2401.03892,
title = {Sampling in Unit Time with Kernel Fisher-Rao Flow},
author = {Aimee Maurais and Youssef Marzouk},
journal= {arXiv preprint arXiv:2401.03892},
year = {2024}
}
Comments
To appear at ICML 2024. Updated with additional numerical examples