Source-linked AI summary
Neural Sampling with Reweighted Normalizing Flows via the Wasserstein--Fisher--Rao JKO Scheme
Chenguang Duan, Johannes Hertrich, Gabriele Steidl
TL;DR
Sampling unnormalized Boltzmann densities is challenging, especially for high-dimensional targets with separated modes. The paper develops a neural WFR JKO sampler that combines transport and reaction, and proves geometric KL convergence for exact iterates without requiring log-concavity or a logarithmic Sobolev inequality. In experiments, its sampler is reported among the best-performing comparisons, while its runtime avoids the exponential dependence on resampling steps reported for neural JKO IC.
Problem
High-dimensional Boltzmann targets with well-separated modes are difficult to sample because accurate sampling must explore relevant regions and recover their relative probability masses.
Source: §1 · p. 1Method
The method minimizes extended KL over finite measures with WFR JKO updates, implemented using neural transport and reaction components in a computable normalizing-flow formulation.
Source: §1 · p. 3, Contributions · p. 3, §3.1 · p. 7Results
Exact WFR JKO iterates converge geometrically in KL with a contraction factor depending only on step size and WFR parameter, and the sampler is reported best or second best among comparisons.
Source: §1 · p. 3, §6.4 · p. 19Takeaways & Limitations
The variational updates use the unnormalized target without requiring its normalizing constant, while approximating both target shape and total mass.
Source: §1 · p. 3The WFR JKO scheme relies on densities: finite KL iterates are absolutely continuous with respect to the target and Lebesgue measure.
Source: §3.1 · p. 8
Abstract
from arXiv · showhide
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.
1 Introduction
Sampling unnormalized Boltzmann densities is difficult for high-dimensional, multimodal targets because it requires exploring modes and recovering their relative masses. This paper combines WFR JKO updates with neural transport and reaction, proving target-independent geometric KL decay and evaluating a practical sampler with resampling and Metropolis–Hastings rejuvenation.
1 Introduction
- WFR combines spatial transport with mass creation or removal, allowing separated-mode masses to change without moving all redistributed mass through low-density regions.The paper motivates this geometry because Wasserstein transport can be metastable for multimodal targets, while convergence estimates often rely on an LSI whose constants can worsen as modes separate.Source: §1 · p. 2
- The target is an unnormalized Boltzmann density, and the method needs neither prior knowledge of its normalizing constant nor a separate normalization for its variational updates.The extended KL objective is minimized at the unnormalized target measure, so the scheme approximates both target shape and total mass.Source: §1 · pp. 1, 3
- For each fixed time step, exact WFR JKO iterates have geometric KL decay with a contraction factor depending only on the step size and WFR parameter, not target convexity or an LSI.The result applies to strictly positive, integrable targets and requires neither log-concavity nor a logarithmic Sobolev inequality.Source: Contributions · p. 3, §1 · p. 3
- The method parameterizes transport and reaction with neural networks, yielding a computable WFR JKO objective that tracks density and total mass without kernel density estimation.Its characteristic representation combines a continuous normalizing flow with a learned positive reweighting function.Source: Contributions · p. 3, §1 · p. 3
- The practical sampler pairs learned transport–reaction updates with resampling and Metropolis–Hastings rejuvenation, using proposals in both data and latent spaces to mitigate particle degeneracy.The experiments compare the method with established MCMC and neural sampling approaches on multimodal targets.Source: Contributions · p. 3, §1 · p. 3
2 Preliminaries
The preliminaries introduce the measure spaces and transport and reaction geometries used to formulate the paper’s WFR JKO objective. They describe Wasserstein motion, Fisher–Rao mass change, and their combination in WFR.
2 Preliminaries
- The paper works with finite positive measures, including probability measures and those with finite second moments, and states that KL is nonnegative and vanishes only for equal measures.KL sublevel sets are weakly sequentially compact.Source: §2 · p. 4
2.1 Wasserstein Metric.
- The Wasserstein-2 geometry describes transport between probability measures, with the Benamou–Brenier formulation expressing distance through the minimum kinetic energy of a connecting path.The section relates absolutely continuous measure curves to solutions of the continuity equation.Source: §2.1 · p. 5
The curve
- In the Fisher–Rao geometry, measure curves evolve through a reaction equation, whose pointwise rate specifies mass creation.The dynamic representation supplies the reaction component later combined with transport in WFR.Source: The curve · p. 5
2.3 Wasserstein–Fisher–Rao Metric.
- WFR combines transport and reaction on finite measures, interpolates between Wasserstein and Fisher–Rao geometries, and metrizes weak convergence.The dynamic formulation uses transport velocity and reaction coefficients subject to Hellinger–Kantorovich regularity conditions.Source: §2.3 · p. 6
3 Wasserstein–Fisher–Rao JKO Scheme
The section formulates the WFR JKO scheme for minimizing the KL functional and proves exponential convergence of its exact iterates toward the target. It studies convergence of the discrete sequence as the iteration count grows.
3.1 WFR JKO Scheme.
- Each WFR JKO step is an implicit gradient-descent update in the WFR metric, defined by minimizing an objective over positive measures.The scheme is considered here for minimizing the KL functional.Source: §3.1 · p. 7, §3 · p. 7
- A minimizer for a WFR JKO step is attained when the current measure has finite KL divergence with respect to π.Source: §3.1 · p. 7
3.2 Convergence of WFR JKO for k
- The WFR JKO iterates converge exponentially to π as k tends to infinity.The section studies this discrete-iteration limit, in contrast to analyses of continuous-time limits as the step size tends to zero.Source: §3.2 · pp. 8, 11, §3.1 · p. 8
- The convergence proof uses a descent lemma that bounds the WFR metric by the Fisher–Rao metric and applies the iterate's optimality condition.Source: §3.2 · p. 8
4 Normalizing Flow via WFR JKO Scheme
The section develops a computable loss for a WFR JKO step and uses neural networks to learn its transport velocity and reaction rate. The algorithm repeats this learning and flow-update procedure, then resamples weighted samples to form the next ensemble.
4.1 Loss Function. We consider the admissible class
- The section reformulates the WFR JKO objective into a computable loss for one step of the scheme.Proposition 4.1 reformulates the action term, and the authors suggest an objective based on that proposition.Source: §4.1 · pp. 12, 14
4.2 Algorithms.
- The algorithm learns neural velocity and reaction functions by minimizing the WFR JKO step loss at each iteration.The learned functions are parameterized as neural networks.Source: §4.2 · p. 14
- After training, the algorithm solves the flow ODE and applies the reaction weights to obtain weighted samples and the next measure's total mass.Source: §4.2 · p. 14
- The algorithm resamples the weighted samples to obtain an unweighted ensemble for the next iteration.Source: §4.2 · p. 14
5 Metropolis–Hastings Steps
The section uses Metropolis–Hastings steps after WFR JKO resampling to address particle degeneracy while preserving the current normalized measure. It describes Gaussian, Langevin, and latent-space proposals, and a flow-based procedure for evaluating intermediate densities.
5 Metropolis–Hastings Steps
- After resampling, a Metropolis–Hastings step rejuvenates the ensemble while preserving its distribution, since the current normalized measure is stationary for the transition kernel.Resampling discards low-weight particles and duplicates high-weight ones, collapsing the ensemble's support.Source: §5.1 · p. 15, §5 · p. 15
5.2 Application to the WFR–JKO Scheme.
- The section considers Gaussian, Langevin, and latent-space proposals for Metropolis–Hastings steps within WFR JKO sampling.For the Langevin proposal, the target gradient is replaced by ∇V because it is directly computable and ρk is expected to be close to π for large k.Source: §5.2 · p. 16
- The latent-space proposal maps samples back through the flow, applies a base Markov kernel there, and maps proposals back to data space.The base kernel may be Gaussian or Langevin, with Gaussian energy for the latter.Source: §5.2 · p. 16
5.3 Density Evaluation of ρk.
- To evaluate ρk at arbitrary points, the section concatenates the stepwise velocities and reactions and uses a reversed flow from ρk to ρ0.Source: §5.3 · p. 17
6 Numerical Results
The numerical section evaluates the WFR–JKO sampler on challenging test distributions, compares it with MCMC and neural sampling methods, and reports that it is consistently among the best performers.
6.1 Test Distributions.
- The test suite includes Mustache, whose long, narrow tails at σ = 0.9 make it difficult for classical MCMC methods to sample.Its transformation has determinant one, so its normalization constant is one.Source: §6.1 · p. 18
- Shifted 8 Modes uses eight equally weighted Gaussians centered on a circle shifted from the origin, while Shifted 8 Peaky narrows those modes until they are disconnected.The shifted center challenges classical MCMC mass allocation, and the narrower modes make sampling harder.Source: §6.1 · p. 18
6.2 Implementation and Hyperparameters.
- The implementation starts with a small WFR–JKO step size and increases it geometrically, resampling and applying a Gaussian-proposal Metropolis–Hastings step after each iteration.Velocity and reaction fields use a six-block ResNet with two-layer, 512-hidden-neuron SiLU blocks; ODEs use RK4 with step size 0.1.Source: §6.2 · p. 18
6.3 Comparisons.
- Comparisons include MALA, HMC, DDS, CRAFT, neural JKO, and importance-corrected neural JKO.The methods are those considered in.Source: §6.3 · p. 18
6.4 Results.
- Energy distance is computed between 50,000 generated and ground-truth samples, and Figure 1 plots their first two coordinates for four multimodal examples.The figure's rows show shifted 8 modes, shifted 8 peaky, GMM-10, and GMM-20.Source: §6.4 · pp. 18–19, Figure 1 · p. 20
- WFR–JKO is always the best or second-best method in the comparisons, while all other methods degenerate on high-dimensional multimodal targets.Neural JKO IC is the only competitive comparison; unlike that method, WFR–JKO does not have runtime that grows exponentially with the number of resampling steps.Source: §6.4 · p. 19
- Table 2 compares energy distances between generated and ground-truth samples across test distributions; smaller values indicate better results.The table marks the best result for each distribution in green and the second-best in blue.Source: Table 2 · p. 19, §6.4 · p. 19
A Auxiliary Proofs
The auxiliary proof establishes weak sequential compactness of bounded KL sublevel sets by showing uniform tightness and using weak convergence.
A Auxiliary Proofs
- For r > 0 and finite measure π on R^d, the set of measures with KL(ρ∥π) ≤ r is weakly sequentially compact.The proof first establishes uniform tightness and then applies Prokhorov’s theorem and weak lower semicontinuity of KL.Source: A Auxiliary Proofs · p. 26
- The proof uses the convex dual representation x log x − x + 1 = sup_a{ax − e^a + 1} to rewrite the KL divergence.This representation supports the subsequent tightness argument.Source: A Auxiliary Proofs · p. 26
- Tightness follows by choosing a compact set using the finiteness and tightness of π, then bounding mass outside it uniformly over the KL sublevel set.The proof tests with a scaled indicator of the compact set’s complement.Source: A Auxiliary Proofs · p. 26
B Pseudo-code of the Algorithms
The appendix provides pseudo-code for empirical-risk evaluation, a WFR–JKO update, and resampling weighted particles into an unweighted ensemble.
B Pseudo-code of the Algorithms
- Algorithm 1 evaluates the empirical risk by sampling Rademacher variables, simulating the ODE system, and computing per-particle losses.Its output is the empirical risk for the current particles, velocity, reaction, and log-mass.Source: B Pseudo-code of the Algorithms · p. 27
- Algorithm 2 applies a JKO step by simulating the ODE system for the particles and returning updated weighted particles.The procedure takes the current particles, velocity, and reaction as inputs.Source: B Pseudo-code of the Algorithms · p. 27
- The resampling procedure normalizes particle weights, samples particle indices from the resulting multinomial distribution, and returns unweighted particles.The JKO-step output includes updated samples, weights, log-density, and total mass.Source: B Pseudo-code of the Algorithms · pp. 27–28