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

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

    Sep 16, 2026Francesca Romana Crucinio, Sahani PathirajaWasserstein Gradient FlowsLog-Concave Distributions

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

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

    May 5, 2026Aaron Havens, Brian Karrer, Neta ShaulDiffusion SamplingRiemannian Manifolds

    FAIR at Meta · Weizmann Institute of Science

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

    Oct 1, 2026Daniel McBride, Pratik Khandagale, Cristina Garcia-Cardona +1Wasserstein Gradient FlowsLatent Transport

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