Weighted Stochastic Differential Equation to Implement Wasserstein-Fisher-Rao Gradient Flow

TL;DR

This paper introduces a weighted stochastic differential equation framework based on Wasserstein-Fisher-Rao geometry to improve sampling in complex, non-log-concave distributions, providing a new theoretical foundation for advanced generative modeling.

Contribution

It formulates WFR-based sampling dynamics using explicit correction terms and weighted SDEs, offering a rigorous geometric and operator-theoretic analysis as a foundation for future work.

Findings

Provides a novel weighted SDE formulation for WFR gradient flows

Clarifies the geometric and operator structure of WFR-based sampling

Lays groundwork for improved sampling algorithms in complex distributions

Abstract

Score-based diffusion models currently constitute the state of the art in continuous generative modeling. These methods are typically formulated via overdamped or underdamped Ornstein--Uhlenbeck-type stochastic differential equations, in which sampling is driven by a combination of deterministic drift and Brownian diffusion, resulting in continuous particle trajectories in the ambient space. While such dynamics enjoy exponential convergence guarantees for strongly log-concave target distributions, it is well known that their mixing rates deteriorate exponentially in the presence of nonconvex or multimodal landscapes, such as double-well potentials. Since many practical generative modeling tasks involve highly non-log-concave target distributions, considerable recent effort has been devoted to developing sampling schemes that improve exploration beyond classical diffusion dynamics. A promising line of work leverages tools from information geometry to augment diffusion-based samplers with controlled mass reweighting mechanisms. This perspective leads naturally to Wasserstein--Fisher--Rao (WFR) geometries, which couple transport in the sample space with vertical (reaction) dynamics on the space of probability measures. In this work, we formulate such reweighting mechanisms through the introduction of explicit correction terms and show how they can be implemented via weighted stochastic differential equations using the Feynman--Kac representation. Our study provides a preliminary but rigorous investigation of WFR-based sampling dynamics, and aims to clarify their geometric and operator-theoretic structure as a foundation for future theoretical and algorithmic developments.