math.NAOct 7, 2026

Neural Sampling with Reweighted Normalizing Flows via the Wasserstein--Fisher--Rao JKO Scheme

Authors: Chenguang Duan, Johannes Hertrich, Gabriele Steidl

Organizations: Institut für Geometrie und Praktische Mathematik, RWTH Aachen University · Institute of Computer Science, University of Göttingen · Institute of Mathematics, TU Berlin

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

We propose a neural algorithm for sampling from distributions specified by unnormalized Boltzmann densities. Our approach is based on the Jordan--Kinderlehrer--Otto scheme for the Kullback--Leibler divergence in the Wasserstein--Fisher--Rao geometry (WFR JKO scheme). Our contributions are twofold. First, we prove that, for any fixed step size, the exact WFR JKO iterates converge exponentially fast to the target as the number of iterations tends to infinity. Notably, this result requires no structural assumptions on the target, such as log-concavity or a logarithmic Sobolev inequality. Second, we develop a neural implementation of the WFR JKO scheme that parametrizes its transport and reaction components using reweighted normalizing flows. Numerical experiments on challenging multimodal targets demonstrate the promising performance of the proposed method.

Figures & tables

Distributionddh0h_{0}τ\tauα\alphaMH step sizeKKσ0\sigma_{0}
Mustache20.00254.01.00.03201.5
shifted 8 modes20.00254.01.00.03201.5
shifted 8 peaky20.00254.01.00.03201.5
GMM-10100.00254.01.00.03201.5
GMM-20200.00254.01.00.03101.5
GMM-50500.00254.01.00.03101.5
Table 1: Hyperparameters for the WFR JKO algorithm: dd is the dimension of the distribution, h0h_{0} is the initial step size in the WFR–JKO algorithm, τ\tau is the growth of this step size, α\alpha is the parameter from the WFR metric, KK the number of WFR–JKO steps used, and σ0\sigma_{0} is the scaling of the initial distribution. In addition, we report the step sizes for solving the ODEs and for generating the proposals in the MH step.
DistributionMALAHMCDDSCRAFTNeural JKONeural JKO ICWFR JKO (ours)Sampling Error
Mustache4.6×10−24.6\times 10^{-2}1.7×10−21.7\times 10^{-2}6.9×10−26.9\times 10^{-2}9.2×10−29.2\times 10^{-2}1.8×10−21.8\times 10^{-2}\pagecolorgreen!15!white2.9×10−3\pagecolor{green!15!white}\bf 2.9\times 10^{-3}\pagecolorblue!15!white1.2×10−2\pagecolor{blue!15!white}\bf 1.2\times 10^{-2}8.6×10−58.6\times 10^{-5}
shifted 8 modes5.3×10−35.3\times 10^{-3}4.1×10−54.1\times 10^{-5}1.2×10−21.2\times 10^{-2}5.2×10−25.2\times 10^{-2}1.3×10−11.3\times 10^{-1}\pagecolorgreen!15!white1.2×10−5\pagecolor{green!15!white}\bf 1.2\times 10^{-5}\pagecolorblue!15!white1.6×10−5\pagecolor{blue!15!white}\bf 1.6\times 10^{-5}2.6×10−52.6\times 10^{-5}
shifted 8 peaky1.3×10−11.3\times 10^{-1}1.2×10−11.2\times 10^{-1}1.1×10−21.1\times 10^{-2}5.2×10−25.2\times 10^{-2}1.3×10−11.3\times 10^{-1}\pagecolorblue!15!white3.4×10−5\pagecolor{blue!15!white}\bf 3.4\times 10^{-5}\pagecolorgreen!15!white3.1×10−5\pagecolor{green!15!white}\bf 3.1\times 10^{-5}2.4×10−52.4\times 10^{-5}
GMM-101.2×10−21.2\times 10^{-2}1.2×10−21.2\times 10^{-2}3.7×10−33.7\times 10^{-3}1.8×10−11.8\times 10^{-1}1.1×10−21.1\times 10^{-2}\pagecolorblue!15!white5.3×10−5\pagecolor{blue!15!white}\bf 5.3\times 10^{-5}\pagecolorgreen!15!white3.2×10−5\pagecolor{green!15!white}\bf 3.2\times 10^{-5}4.6×10−54.6\times 10^{-5}
GMM-209.1×10−39.1\times 10^{-3}9.1×10−39.1\times 10^{-3}5.0×10−35.0\times 10^{-3}5.4×10−15.4\times 10^{-1}1.0×10−21.0\times 10^{-2}\pagecolorblue!15!white1.1×10−4\pagecolor{blue!15!white}\bf 1.1\times 10^{-4}\pagecolorgreen!15!white1.0×10−4\pagecolor{green!15!white}\bf 1.0\times 10^{-4}6.4×10−56.4\times 10^{-5}
GMM-502.4×10−22.4\times 10^{-2}2.4×10−22.4\times 10^{-2}2.3×10−22.3\times 10^{-2}1.8×1001.8\times 10^{0}2.7×10−22.7\times 10^{-2}\pagecolorgreen!15!white1.0×10−4\pagecolor{green!15!white}\bf 1.0\times 10^{-4}\pagecolorblue!15!white4.1×10−4\pagecolor{blue!15!white}\bf 4.1\times 10^{-4}1.1×10−41.1\times 10^{-4}
Table 2: Energy distance between generated samples and ground truth samples. A smaller value indicates a better result. The best result for each test distribution is marked in green, the second-best is marked in blue.
Figure 1: Plot of the first two coordinates of generated and ground truth samples for different examples. First line: shifted 8 modes; second line: shifted 8 peaky; third line: GMM-10; fourth line: GMM-20.
Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

 ddtϕt(Xki)\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\phi_{t}(X_{k}^{i})=vt∘ϕt(Xki),ϕ0(Xki)=Xki,\displaystyle=v_{t}\circ\phi_{t}(X_{k}^{i}),\quad\phi_{0}(X_{k}^{i})=X_{k}^{i}, 
 ddtψt(Xki)\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\psi_{t}(X_{k}^{i})=gt∘ϕt(Xki),ψ0(Xki)=log⁡ρk(Rd),wt(Xki)=exp⁡(ψt(Xki)),\displaystyle=g_{t}\circ\phi_{t}(X_{k}^{i}),\quad\psi_{0}(X_{k}^{i})=\log\rho_{k}(\mathbb{R}^{d}),\quad w_{t}(X_{k}^{i})=\exp(\psi_{t}(X_{k}^{i})), 
 ddtat(Xki)\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}a_{t}(X_{k}^{i})=wt(Xki){∥vt∘ϕt(Xki)∥22+1λ(gt∘ϕt(Xki))2},a0(Xki)=0,\displaystyle=w_{t}(X_{k}^{i})\Bigl\{\|v_{t}\circ\phi_{t}(X_{k}^{i})\|_{2}^{2}+\frac{1}{\lambda}(g_{t}\circ\phi_{t}(X_{k}^{i}))^{2}\Bigr\},\quad a_{0}(X_{k}^{i})=0, 
 ddtηt(Xki)\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\eta_{t}(X_{k}^{i})=ξi⊤∇vt∘ϕt(Xki)ξi,ηt(Xki)=0.\displaystyle=\xi_{i}^{\top}\nabla v_{t}\circ\phi_{t}(X_{k}^{i})\xi_{i},\quad\eta_{t}(X_{k}^{i})=0. 
Appendix
Algorithm 1 Evaluate the empirical risk
 ddtϕt(Xki)\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\phi_{t}(X_{k}^{i})=vt∘ϕt(Xki),ϕ0(Xki)=Xki,\displaystyle=v_{t}\circ\phi_{t}(X_{k}^{i}),\quad\phi_{0}(X_{k}^{i})=X_{k}^{i}, 
 ddtψt(Xki)\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\psi_{t}(X_{k}^{i})=gt∘ϕt(Xki),ψ0(Xki)=log⁡ρk(Rd),\displaystyle=g_{t}\circ\phi_{t}(X_{k}^{i}),\quad\psi_{0}(X_{k}^{i})=\log\rho_{k}(\mathbb{R}^{d}), 
 ddtζt(Xki)\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\zeta_{t}(X_{k}^{i})=− div vt∘ϕt(Xki),ζ0(Xki)=log⁡ρk(Xki),\displaystyle=-\,\mathrm{div}\,v_{t}\circ\phi_{t}(X_{k}^{i}),\quad\zeta_{0}(X_{k}^{i})=\log\rho_{k}(X_{k}^{i}), 
Appendix
Algorithm 2 Applying a JKO step

Explore similar work

CardsList
Sep 16, 2026stat.ML

Preservation of Log-Concavity and Convergence of Wasserstein-Fisher-Rao Gradient Flows

We study the convergence of Wasserstein-Fisher-Rao (WFR) gradient flows for sampling from probability distributions known up to a normalisation constant. By combining Wasserstein transport with Fisher-Rao birth-death dynamics, WFR flows balance exploration and selection. These flows have been recognised as a promising mechanism to accelerate convergence beyond Langevin dynamics. We show that for a class of strongly log-concave target distributions satisfying additional curvature conditions, WFR flows preserve strong log-concavity, in contrast to Wasserstein flows which enjoy this property only in the Gaussian setting. Exploiting this result, we derive explicit non-asymptotic convergence rates for the symmetrised Kullback-Leibler divergence, without requiring a warm-start as required in current estimates. In particular, we show that the convergence rate decomposes additively into Wasserstein and Fisher-Rao contributions, thereby confirming a recent conjecture within this setting. These results provide refined convergence guarantees and further develop the theoretical foundations of WFR gradient flows for sampling and Bayesian inference.

Francesca Romana Crucinio, Sahani Pathiraja

ESOMAS, University of Turin, Italy & Collegio Carlo Alberto, Turin, Italy · School of Mathematics & Statistics, UNSW Sydney, Australia

Wasserstein Gradient FlowsLog-Concave DistributionsWasserstein DistanceLangevin DynamicsConvergenceGradient
AskPDF
May 5, 2026cs.LG

Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes

Sampling from unnormalized densities is analogous to the generative modeling problem, but the target distribution is defined by a known energy function instead of data samples. Because evaluating the energy function is often costly, a primary challenge is to learn an efficient sampler. We introduce Flow Sampling, a framework built on diffusion models and flow matching for the data-free setting. Our training objective is conditioned on a noise sample and regresses onto a denoising diffusion drift constructed from the energy function. In contrast, diffusion models' objective is conditioned on a data sample and regresses onto a noising diffusion drift. We utilize the interpolant process to minimize the number of energy function evaluations during training, resulting in an efficient and scalable method for sampling unnormalized densities. Furthermore, our formulation naturally extends to Riemannian manifolds, enabling diffusion-based sampling in geometries beyond Euclidean space. We derive a closed-form formula for the conditional drift on constant curvature manifolds, including hyperspheres and hyperbolic spaces. We evaluate Flow Sampling on synthetic energy benchmarks, small peptides, large-scale amortized molecular conformer generation, and distributions supported on the sphere, demonstrating strong empirical performance.

Aaron Havens, Brian Karrer, Neta Shaul

FAIR at Meta · Weizmann Institute of Science

Diffusion SamplingRiemannian ManifoldsConformational Ensembles
AskPDF
Oct 1, 2026stat.ML

Wasserstein Gradient Flows and Forward-Only Diffusion Are Not Enough for Multimodal Sampling

There has been a proliferation of sampling algorithms based on Wasserstein gradient flows (WGF) and forward-only diffusion processes (FODP), often accompanied by theoretical guarantees of exponentially fast convergence to the target distribution. These guarantees are frequently interpreted as evidence that such methods can efficiently sample complex multimodal distributions, often supported by empirical results. In this work, we argue that this interpretation is fundamentally misleading. By invoking the Jordan-Kinderlehrer-Otto (JKO) scheme and Otto calculus, we establish that the canonical WGF sampling dynamics and overdamped forward diffusion share the same density evolution and therefore inherit the same metastability and slow-mixing phenomena long understood in nonequilibrium statistical physics. We analyze this family of samplers using two complementary tools -- spectral analysis and mean first-passage time (MFPT) analysis -- and show that well-separated multimodality can induce exponentially long mixing times associated with small spectral gaps and rare inter-mode transitions. For the commonly adopted log-linear annealing schedule studied here, we find that introducing intermediate distributions does not remove the exponential scaling of the total transport time. The limitation is structural rather than implementation-specific: purely local, gradient-driven transport mechanisms can require exponentially long times to transport probability mass across well-separated modes. We argue that this represents a fundamental limitation of WGF- and FODP-based sampling in their standard forms, and motivates future development of fundamentally nonlocal mechanisms for efficient multimodal sampling.

Daniel McBride, Pratik Khandagale, Cristina Garcia-Cardona +1

Los Alamos National Laboratory, Los Alamos, NM 87545, USA

Wasserstein Gradient FlowsLatent TransportDiffusion ModelsMultimodal Generation
AskPDF