Wasserstein-Fisher-Rao Gradient Flow

• Wasserstein-Fisher-Rao gradient flow is a geometric framework unifying optimal transport and birth–death dynamics by coupling L2-Wasserstein and Fisher–Rao metrics.
• It utilizes gradient descent on free energy functionals with operator splitting and JKO discretization to achieve exponential minimization of divergence measures.
• Practical applications include advanced sampling, multi-objective optimization, and mixture model learning, offering robust performance in high-dimensional probabilistic modeling.

The Wasserstein-Fisher-Rao (WFR) gradient flow constitutes a geometric framework for the evolution of probability measures, combining optimal transport and birth-death dynamics within a single Riemannian structure. WFR interpolates smoothly between the $L^2$–Wasserstein metric, which governs mass transport, and the Fisher–Rao metric, which encodes growth and decay phenomena. In contrast to classical Wasserstein flows that conserve total mass, WFR allows mass creation and annihilation, enabling more flexible and robust algorithms for sampling, optimization, and learning in high-dimensional probabilistic spaces. Its gradient flow—interpreted as the steepest descent of functionals such as the Kullback–Leibler divergence—forms the basis of a rapidly growing body of work in contemporary computational statistics, generative modeling, and multi-objective optimization.

1. Geometry and Mathematical Formalism

The WFR metric is defined on the space of positive probability densities $\mathcal{P}_{ac}(\mathbb{R}^d)$ via the dynamic Benamou–Brenier formulation:

$$d_{\mathrm{WFR}}^2(\rho_0, \rho_1) = \inf_{(\rho_t, v_t, r_t)} \int_0^1 \int_{\mathbb{R}^d} \left( |v_t(x)|^2 + |r_t(x)|^2 \right) \rho_t(x) dx dt$$

subject to the continuity-reaction equation:

$$\partial_t \rho_t + \nabla \cdot (\rho_t v_t) = r_t \rho_t, \quad \rho_0 = \rho_0, \; \rho_1 = \rho_1.$$

A tangent vector at $\mu$ admits a canonical decomposition:

$$\partial_t \mu = -\nabla \cdot (\mu \nabla \phi) + \mu \psi,$$

where $(\phi,\psi)$ are potentials and reaction rates, with the induced inner product:

$$\langle (\phi, \psi), (\phi', \psi') \rangle_{WFR, \mu} = \int \mu(x)(|\nabla \phi(x)|^2 + \psi(x)^2) dx.$$

The pure-Wasserstein ($\psi \equiv 0$) case corresponds to advective transport, and pure-Fisher–Rao ($\phi \equiv 0$) to birth–death mechanisms (Crucinio et al., 6 Jun 2025, Crucinio et al., 22 Nov 2025, Yan et al., 2023).

2. WFR Gradient Flows for KL Minimization

WFR flows are typically formulated as gradient flows of free energy functionals. For minimization of reverse KL, $D_{KL}(\mu \| \pi)$, the evolution is governed by:

$$\partial_t \mu = \nabla \cdot [\mu \nabla \log(\mu/\pi)] + \mu [\log(\pi/\mu) - \mathbb{E}_\mu[\log(\pi/\mu)]],$$

or, equivalently, $\partial_t \mu = f_W(\mu) + f_{FR}(\mu)$, with each component representing transport and birth–death, respectively (Crucinio et al., 6 Jun 2025, Crucinio et al., 22 Nov 2025).

For inclusive KL minimization, i.e., minimizing $\mathrm{KL}(\pi \| \mu)$, the WFR gradient flow yields:

$$\partial_t \mu_t = \alpha\,\nabla \cdot \Big( \mu_t \nabla [1 - \tfrac{d\pi}{d\mu_t}] \Big) - \beta\,\mu_t \Big( 1 - \tfrac{d\pi}{d\mu_t} \Big),$$

with exponential convergence rate for $\mathrm{KL}(\pi \| \mu_t)$ under smoothness and moment bounds (Zhu, 31 Oct 2024).

3. Discretizations: JKO Scheme and Operator Splitting

The Jordan-Kinderlehrer-Otto (JKO) implicit Euler discretization for WFR is given by:

$$\mu_{k+1} = \arg\min_{\mu} \left\{ \frac{1}{2\tau} d_{\mathrm{WFR}}^2 (\mu, \mu_k) + \mathrm{KL}(\pi \| \mu) \right\}$$

where the Euler–Lagrange condition involves the discrete WFR gradient augmented by the functional's first variation.

Operator splitting schemes (Lie–Trotter) numerically approximate the coupled WFR flow by alternating pure-Wasserstein and pure-Fisher–Rao steps. Two orders are defined:

• W–FR splitting (Wasserstein followed by Fisher–Rao): $$S_W(\tau, \nu_k) \rightarrow S_{FR}(\tau, \widehat{\nu}_{k+\frac{1}{2}})$$
• FR–W splitting (Fisher–Rao then Wasserstein): $$S_{FR}(\tau, \eta_k) \rightarrow S_W(\tau, \widehat{\eta}_{k+\frac{1}{2}})$$

Closed-form Gaussian solutions demonstrate that operator ordering, step size, and initial covariance determine convergence speed; in some scenarios, suitably chosen W–FR splitting outpaces even the exact WFR flow (Crucinio et al., 22 Nov 2025). These schemes are $O(\tau)$ accurate but introduce splitting-induced commutator corrections.

4. Sequential Monte Carlo and Particle Approximations

Particle-based discretizations provide practical algorithms for WFR flows.

The SMC–WFR algorithm applies:

1. W-mutation: Move particles via ULA:

$$X_{n-\frac{1}{2}}^i = X_{n-1}^i + \gamma \nabla \log \pi(X_{n-1}^i) + \sqrt{2\gamma} \xi^i$$

2. FR-step reweighting: Assign importance weights using the empirical density and target:

$$w_n^i \propto \left( \frac{\pi(X_{n-1/2}^i)}{\hat{\mu}_{n-1/2}(X_{n-1/2}^i)} \right)^{\delta},\quad \delta = 1 - e^{-\gamma}$$

3. Resample: Produce new cloud with uniform weights.

Under a log-Sobolev assumption, discrete-time convergence is exponential up to $O(\gamma)$ errors. Empirical comparisons yield significantly lower MMD, decreased bias in sample moments, and accelerated mixing relative to pure Langevin and birth–death schemes (Crucinio et al., 6 Jun 2025, Ren et al., 2023).

5. Theoretical Properties: Sharp Decay, Log-Concavity, and Uniqueness

The exact WFR flow for $\mathrm{KL}(\mu \| \pi)$ preserves strong log-concavity: if $\mu_0 \propto e^{-V_0}$ and $\pi \propto e^{-V_\pi}$ are strongly log-concave, then the evolved density remains so for all $t > 0$. For birth–death-only (FR), log-concavity is interpolated via the explicit semigroup; for pure-Wasserstein, a Prékopa–Leindler argument bounds short-time preservation; for coupled flows, convexity restoration by FR yields uniform bounds (Crucinio et al., 22 Nov 2025).

Sharp exponential dissipation of free energy is established for KL and Jeffrey's divergence:

$$J(\mu_t, \pi) = \mathrm{KL}(\mu_t \| \pi) + \mathrm{KL}(\pi \| \mu_t)$$

with dissipation rate at least $2 \min(\alpha_\pi, \alpha_t) + 1$.

Existence, uniqueness, and stability of WFR gradient flows follow from metric-space theory, with contractivity in the Fisher–Rao direction (Zhu, 31 Oct 2024).

6. Applications: Multi-Objective Optimization and Mixture Model Learning

WFR flows are employed in multi-objective optimization (MOO), where the gradient flow is given by:

$$\partial_t \rho_t = \nabla \cdot (\rho_t \nabla \delta_\rho E[\rho_t]) - \rho_t (\delta_\rho E[\rho_t] - \int \delta_\rho E[\rho_t] d\rho_t)$$

where $E[\rho]$ combines objective alignment, dominance, repulsion, and entropy terms. Splitting dynamics ensure global Pareto optimality, outperforming repulsive-only MOO methods by relocating dominated particles (Ren et al., 2023).

In learning Gaussian mixtures, WFR gradient descent alternately updates particle weights and positions to minimize a nonparametric negative log-likelihood. The resulting algorithm escapes local minima that trap Wasserstein-only and Fisher–Rao-only schemes, achieving lower training and test error and near-zero sub-optimality gap in empirical studies (Yan et al., 2023).

7. Limitations, Tempered Flows, and Algorithmic Extensions

Tempering—replacing the target with a geometric mixture $\pi_t \propto \pi^t \mu_0^{1-t}$—does not accelerate convergence; in practice, tempered-WFR flows converge more slowly or at best equally rapidly as untempered flows, with explicit upper bounds on the KL decay (Crucinio et al., 6 Jun 2025).

Kernelized approximations (MMD and KSD flows) regularize transport forces for feasible high-dimensional implementation, inheriting the asymptotic properties of the underlying WFR structure (Zhu, 31 Oct 2024).

Potential extensions include entropic regularized numerics, deep JKO networks, mirror-descent analysis, and further paper of the $\alpha, \beta$ trade-off in balancing transport and birth–death mechanisms. Outstanding questions remain in high-dimensional guarantees, optimal splitting, and rigorous convergence of kernelized flows as regularization bandwidth diminishes (Zhu, 31 Oct 2024).

---

The WFR gradient flow unifies optimal transport and birth–death dynamics under a canonical geometric structure, enabling exponentially fast minimization of divergence functionals, practical sampling, and robust optimization in probabilistic modeling. Its algorithmic discretizations are provably consistent and empirically superior to schemes using only transport or reweighting, with ongoing developments in theory and applications spanning statistics, machine learning, and multi-objective optimization.