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

TL;DR

The paper introduces Jeffreys Flow, a new robust generative method for sampling complex physical systems that overcomes mode collapse by using symmetric divergence and distillation from Parallel Tempering data.

Contribution

It proposes Jeffreys Flow, which mitigates mode collapse in Boltzmann generators by employing the symmetric Jeffreys divergence and empirical data distillation.

Findings

Reduces mode collapse in multi-modal distributions.

Improves accuracy and scalability on complex benchmarks.

Accelerates importance sampling in quantum thermal state simulations.

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.