Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization

TL;DR

This paper introduces off-policy log-dispersion regularization (LDR), a novel method that improves the training efficiency of Boltzmann generators by leveraging additional energy information, leading to better performance with less data.

Contribution

The paper proposes LDR, a new regularization framework that enhances Boltzmann generator training by utilizing off-policy data and target energy labels, without requiring on-policy samples.

Findings

LDR improves performance across benchmarks.

Sample efficiency increases up to tenfold.

Supports various datasets and training methods.

Abstract

Sampling from unnormalized probability densities is a central challenge in computational science. Boltzmann generators are generative models that enable independent sampling from the Boltzmann distribution of physical systems at a given temperature. However, their practical success depends on data-efficient training, as both simulation data and target energy evaluations are costly. To this end, we propose off-policy log-dispersion regularization (LDR), a novel regularization framework that builds on a generalization of the log-variance objective. We apply LDR in the off-policy setting in combination with standard data-based training objectives, without requiring additional on-policy samples. LDR acts as a shape regularizer of the energy landscape by leveraging additional information in the form of target energy labels. The proposed regularization framework is broadly applicable, supporting unbiased or biased simulation datasets as well as purely variational training without access to target samples. Across all benchmarks, LDR improves both final performance and data efficiency, with sample efficiency gains of up to one order of magnitude.