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

2026-04-06

AI Summary

Sampling physical systems on rugged energy landscapes is often hindered by rare events and metastable traps. Conventional Boltzmann generators, trained with reverse Kullback–Leibler (KL) divergence, suffer from mode collapse and struggle to capture multimodal distributions. This work proposes the Jeffreys Flow framework, which for the first time incorporates the symmetric Jeffreys divergence into generator training. By distilling empirical data from parallel tempering trajectories, the method enhances global multimodal coverage while preserving local accuracy. It effectively mitigates mode collapse and structurally corrects generative bias, demonstrating superior scalability and accuracy on high-dimensional non-convex benchmark tasks. Furthermore, it successfully rectifies the gradient bias in Replica Exchange Stochastic Gradient Langevin Dynamics, significantly accelerating importance sampling in path-integral Monte Carlo simulations of quantum thermal states.

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.

Problem

• rare event sampling
• mode collapse
• multi-modal distributions
• energy landscapes
• Boltzmann generators

Innovation

• Jeffreys divergence
• Boltzmann generators
• Parallel Tempering
• mode collapse
• importance sampling