Jeffreys Flow: Robust Boltzmann Generators for Rare Event Sampling via Parallel Tempering Distillation

Guang Lin; Christian Moya; Di Qi; Xuda Ye

cs.LGApr 7, 2026

Jeffreys Flow: Robust Boltzmann Generators for Rare Event Sampling via Parallel Tempering Distillation

Authors: Guang Lin, Christian Moya, Di Qi, Xuda Ye

Organizations: Department of Mathematics, Purdue University, West Lafayette, IN 47907, USA

Abstract

Sampling physical systems with rough energy landscapes is hindered by rare events and metastable trapping. While Boltzmann generators already offer a solution, their reliance on the reverse Kullback--Leibler divergence frequently induces catastrophic mode collapse, missing specific modes in multi-modal distributions. Here, we introduce the Jeffreys Flow, a robust generative framework that mitigates this failure by distilling empirical sampling data from Parallel Tempering trajectories using the symmetric Jeffreys divergence. This formulation effectively balances local target-seeking precision with global modes coverage. We show that minimizing Jeffreys divergence suppresses mode collapse and structurally corrects inherent inaccuracies via distillation of the empirical reference data. We demonstrate the framework's scalability and accuracy on highly non-convex multidimensional benchmarks, including the systematic correction of stochastic gradient biases in Replica Exchange Stochastic Gradient Langevin Dynamics and the massive acceleration of exact importance sampling in Path Integral Monte Carlo for quantum thermal states.

Figures & tables

Figure 1: The overall architecture of the Jeffreys Flow. PT acts as the principal guide across the temperature ladder, providing reference samples (

\mu_{1},\dots,\mu_{M}

) to train the sequence of normalizing flows (

F_{1},F_{2},\dots,F_{M}

). Simultaneously, intermediate states (

\nu_{0},\nu_{1},\dots,\nu_{M}

) are progressively pushed forward and reweighted, intrinsically averting mode collapse without relying solely on the target potential function.

(F_{k},w_{k})=\mathtt{JeffreysFlow}(\mu_{k-1}^{\prime},\mu_{k};U_{k-1},U_{k}).

Algorithm 1 Jeffreys Flow: Sequential Distillation from Parallel Tempering

Panel of Figure 2 — Figure 2: Density plots of the generated samples across four 2D landscapes. The Jeffreys Flow ( $0<\theta<1$ ) robustly corrects the noisy artifacts of Forward KL ( $\theta=0$ ) and strictly prevents the destructive mode collapse inherent to Reverse KL ( $\theta=1$ ).

Figure 3: Left: The generated distribution from the Jeffreys Flow. The scatter samples accurately capture all modes of the periodic potential ( $N=4^{12}$ ). Right: As sample size $N$ increases, the ESS remains consistently optimal ( $\approx 100\%$ ) while the $L^{2}$ bias strictly follows the theoretical Monte Carlo convergence rate.

Figure 4: Sequential distillation on the 3D Gaussian Mixture Model. Top: Reference PT samples at various temperatures. Bottom: Generated samples from the distilled Jeffreys Flow. The resulting bias of the Jeffreys Flow is much smaller than PT.

Figure 5: Sequential distillation on the 4D Rosenbrock. Top: Reference PT samples. Bottom: Generated samples from the Jeffreys Flow successfully learning the highly nonlinear correlations.

Figure 6: Comparison for 8D Nonlinear Rastrigin. Left: Potential energy distribution

P(U)

comparison between PT and Jeffreys Flow. Right: 1D marginal distribution

p(x_{1})

comparison. Both methods show consistent mode coverage and energy profiles.

Step	1	2	3	4	5	6	7	8	9	10	11
$\beta$	0.020	0.032	0.050	0.100	0.150	0.200	0.250	0.300	0.350	0.400	0.450
CESS	93.7%	95.9%	91.6%	77.6%	70.7%	78.1%	81.4%	87.7%	89.9%	92.3%	94.1%
ESS	93.7%	88.5%	78.1%	57.1%	100%	78.1%	57.0%	100%	89.9%	80.1%	71.5%
Step	12	13	14	15	16	17	18	19	20	21	22
$\beta$	0.500	0.550	0.600	0.650	0.700	0.750	0.800	0.850	0.900	0.950	1.000
CESS	94.6%	95.4%	95.4%	95.9%	95.7%	96.1%	95.5%	95.9%	95.0%	95.3%	95.0%

Table 1: Performance metrics for the 8D Nonlinear Rastrigin potential. The table displays the inverse temperature

\beta

, the CESS, and the ESS at each stage. The ESS is reset to

100\%

(marked as boxed) whenever adaptive resampling is performed.

Figure 7: Decorrelation in the 16D Solvated Periodic Grid. Top: PT reference samples exhibit spurious diagonal correlations. Bottom: Jeffreys Flow successfully uncouples the bath to recover the exact independent checkerboard structure.

Figure 8: Potential of Mean Force (PMF) along

x_{1}

. Jeffreys Flow perfectly reconstructs the analytical free energy barriers that PT completely fails to cross.

Figure 9: Decorrelation in 2D GMM. Visual comparison between reSGLD and Jeffreys Flow. Jeffreys (SVRG) fully resolves the probability density distribution. Jeffreys (No SVRG) generates samples with excessive spread and undesirable artificial tails.

Step	0	1	2	3	4	5
$\beta$	0.0	0.05	0.2	0.4	0.7	1.0
reSGLD	5.97e-02	5.90e-02	6.11e-02	7.49e-02	1.12e-01	1.50e-01
Jeffreys (No SVRG)	2.81e-03	2.57e-02	1.91e-02	1.96e-02	1.32e-02	1.48e-02
Jeffreys (SVRG)	4.94e-03	1.65e-03	1.56e-02	6.44e-03	4.60e-03	8.45e-03

Table 2: Bias Evaluation across the intermediate targets in the 2D GMM experiment. The bias is drastically reduced by training with Jeffreys Flow, filtering out aggressive discretization errors inherent in reSGLD.

Step	0	1	2	3	4	5
$\beta$	0.0	0.05	0.2	0.4	0.7	1.0
Jeffreys (No SVRG)	–	99.17%	97.52%	94.42%	87.17%	76.80%
Jeffreys (SVRG)	–	99.25%	96.77%	95.34%	91.77%	86.02%

Table 3: ESS across the intermediate targets in the 2D GMM experiment. High quantitative ESS values demonstrate the ability of Jeffreys Flow as a robust feed-forward Boltzmann generator for direct, independent sample generation.

Step	1	2	3	4	5	6	7	8
Exact	83.93%	91.49%	92.48%	88.69%	77.30%	85.23%	89.56%	89.24%
Full Potential	83.42%	88.28%	93.09%	88.83%	79.90%	86.18%	89.25%	91.15%
Stochastic	61.71%	64.27%	84.69%	90.00%	89.08%	93.23%	79.66%	94.14%

Table 4: CESS across the intermediate targets in the 2D Screened Poisson experiment. All evaluated configurations sustain extremely high CESS scores, confirming the robustness and massive acceleration achieved by the trained Jeffreys Flow.

Figure 10: 2D Screened Poisson Equipment Setup. Left: The true multi-source distribution field generated by the ground truth locations. Right: The resulting physical response $u(x)$ modeled by solving the discretized Poisson PDE. The $81$ fixed local sensors (square markers) collect the noisy measurements used to infer the unknown source locations (star markers).

Figure 11: Marginal Posterior for Source 1 ( $\theta_{1}$ ) in the Poisson Inverse Problem. The Exact configuration produces tightly bounded, high-fidelity posterior samples. The Neural Network-approximated Stochastic result exhibits structural degradation, validating the necessity of using exact PDE evaluations for calculating the importance weights.

Discretization $N$	8	12	16	20	24	28	32
$L^{2}$ Bias	1.60e-02	7.45e-03	5.93e-03	3.56e-03	4.02e-03	2.00e-03	3.84e-04

Table 5: Sampling Bias Evaluation in High Dimensions. Using solely a single transport map trained exclusively at the truncated dimension

N_{0}=8

, the reweighted Jeffreys Flow generates high-fidelity sampling accuracy uniformly up to

N=32

, with the bias adhering to the rapid theoretical

\mathcal{O}(1/N^{2})

discretization error decay.

Figure 12: Quantum Path Integral Distillation. Comparison of the marginal spatial density between the classical distribution constructed from the empirical training samples and the distilled quantum distribution extracted using Jeffreys Flow.

Figure 13: Progression of the CESS over the intermediate bridging steps. Sustained high CESS scores confirm the robustness and stability of the transport map during the classical-to-quantum distillation sequence.

Explore similar work

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows
Sep 25, 2026Louis Grenioux, RuiKang OuYang, Luhuan WuMolecular Generation Normalizing Flows
Center for Computational Mathematics, Flatiron Institute, New York, NY, USA · Department of Engineering, University of Cambridge, Cambridge, UK · Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD, USA
FUND: Density Flow for Sampling Unnormalised Distributions
Sep 26, 2026Vikas Kanaujia, Vipul AroraFlow-Based Generative Modeling Flow Matching
Department of Electrical Engineering IIT Kanpur · Department of Electrical Engineering KU Leuven
Leveraging generative models to assist Monte Carlo sampling
Aug 7, 2026Marylou GabriéNormalizing Flows Generative Modeling
Laboratoire de Physique Statistique, École normale supérieure, PSL Research University, 24 rue Lhomond, 75005 Paris, France

Sep 25, 2026J/K move · Enter open · S save

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows

Louis Grenioux, RuiKang OuYang, Luhuan Wu

Center for Computational Mathematics, Flatiron Institute, New York, NY, USA · Department of Engineering, University of Cambridge, Cambridge, UK · Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD, USA

Molecular Generation Normalizing Flows Flow Map Learning Importance Sampling

Open paper AskPDF arXiv

Generating independent, equilibrium samples of molecular systems at scale remains a central obstacle in computational statistical mechanics. Boltzmann Generators address this by pairing a generative model with importance sampling to obtain consistent samples from the target distribution. We introduce Normalizing Flow Flow Maps (NF

^2

M), which combines the strengths of recent stochastic flow maps with the tractability of classic normalizing flows to build a Boltzmann Generator. Unlike most methods, which correct the generative model only at the end, NF

^2

M reweighs each denoising transition as generation proceeds, avoiding wasted compute on trajectories that are ultimately discarded. At each denoising step, a conditional normalizing flow proposes clean configurations given the current noisy state (a simpler task than sampling directly from the target) and its exact likelihood enables correcting each proposal toward the true denoising transition of the target Boltzmann distribution. This is in contrast to most existing methods, whose likelihoods are approximate or expensive to evaluate, undermining the statistical reliability of the correction. We establish consistency of the corrected transitions and bound how approximation errors propagate through the sampling chain. We evaluate NF

^2

M on peptide systems, demonstrating improved sampling efficiency and sample quality.

Abstract

Explore similar work

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows

FUND: Density Flow for Sampling Unnormalised Distributions

Leveraging generative models to assist Monte Carlo sampling