Table of Contents
Fetching ...

Sampling in Unit Time with Kernel Fisher-Rao Flow

Aimee Maurais, Youssef Marzouk

TL;DR

We address sampling from an unnormalized target $\pi_1$ using a gradient-free, finite-time transport from a reference $\pi_0$ by evolving along the unit-time geometric path $\pi_t \propto \pi_0^{1-t}\pi_1^t$. The core contribution is a mean-field ODE for a velocity field $v_t$ obtained by solving a Poisson equation in an RKHS, yielding a kernelized, gradient-free interacting particle system (KFRFlow) with discrete-time (KFRFlow-I) and stochastic (KFRD) variants that transport samples toward $\pi_1$. The approach unifies a discrete-time sample-driven optimal transport view with a continuous Fisher--Rao gradient flow perspective, and empirical results demonstrate strong performance against gradient-free baselines and competitiveness with gradient-based methods, especially in higher dimensions. This framework provides a scalable, gradient-free tool for challenging sampling tasks under unnormalized densities, with potential extensions to kernel acceleration, dimension reduction, and adaptive time stepping.

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ère 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.

Sampling in Unit Time with Kernel Fisher-Rao Flow

TL;DR

We address sampling from an unnormalized target using a gradient-free, finite-time transport from a reference by evolving along the unit-time geometric path . The core contribution is a mean-field ODE for a velocity field obtained by solving a Poisson equation in an RKHS, yielding a kernelized, gradient-free interacting particle system (KFRFlow) with discrete-time (KFRFlow-I) and stochastic (KFRD) variants that transport samples toward . The approach unifies a discrete-time sample-driven optimal transport view with a continuous Fisher--Rao gradient flow perspective, and empirical results demonstrate strong performance against gradient-free baselines and competitiveness with gradient-based methods, especially in higher dimensions. This framework provides a scalable, gradient-free tool for challenging sampling tasks under unnormalized densities, with potential extensions to kernel acceleration, dimension reduction, and adaptive time stepping.

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ère 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.
Paper Structure (32 sections, 1 theorem, 53 equations, 20 figures, 2 tables)

This paper contains 32 sections, 1 theorem, 53 equations, 20 figures, 2 tables.

Key Result

Theorem 4.1

In the limit $\Delta t \to 0$, eq:IPS_importance approaches eq:IPS_ode. Thus the IPS obtained via sample-driven optimal transport eq:IPS_importance can be viewed as arising from mean-field model eq:meanfield.

Figures (20)

  • Figure 1: We employ a homotopy-based sampling scheme in this work, deriving a mean-field ODE which approximately transports a reference $\pi_0$ to a target $\pi_1$ in unit time. In discrete time this approach amounts to obtaining incremental maps $T_1, \dots, T_N$.
  • Figure 2: Two-dimensional posteriors: samples at $t \in \{0, 0.25, 0.5, 0.75, 1\}$ generated by KFRFlow \ref{['eq:IPS_ode']} for the donut (top), butterfly (middle), and spaceships (bottom) examples.
  • Figure 3: Two-dimensional posteriors: average KSD at stopping time between $\pi_1$ and ensembles of size $J \in \{100, 400\}$ generated by gradient-free samplers. A missing point indicates that a method was unstable at that setting of $N$.
  • Figure 4: Two-dimensional posteriors: average KSD at stopping time between $\pi_1$ and ensembles of size $J \in \{100, 400\}$ generated by gradient-based samplers. A missing point indicates that a method was unstable at that setting of $N$.
  • Figure 5: Two-dimensional posteriors: evolution of KSD with $t$ for the unit-time methods KFRFlow, KFRFlow-I, KFRD, and EKI for ensembles of $J=400$ and $J = 100$ with $\Delta t = 2^{-8}$. KFRD is plotted with dashed lines because it requires gradients, whereas KFRFlow(-I) and EKI are gradient-free.
  • ...and 15 more figures

Theorems & Definitions (1)

  • Theorem 4.1