Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Henrik Schopmans, Christopher von Klitzing, Pascal Friederich

TL;DR

Efficient Training of Boltzmann Generators introduces off-policy log-dispersion regularization (LDR) to improve data efficiency in Boltzmann generator training. By generalizing the log-variance objective to log-dispersion and applying it off-policy as a regularizer on fixed datasets, LDR leverages target energy information without extra on-policy samples. The method supports unbiased, biased, and variational training regimes and is validated across internal-coordinate flows and Cartesian-coordinate TarFlow implementations, showing substantial improvements in final performance and data efficiency, often by an order of magnitude. The work demonstrates broad applicability and practical impact for efficient equilibrium-ensemble sampling in computational chemistry and physics.

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.

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

TL;DR

Efficient Training of Boltzmann Generators introduces off-policy log-dispersion regularization (LDR) to improve data efficiency in Boltzmann generator training. By generalizing the log-variance objective to log-dispersion and applying it off-policy as a regularizer on fixed datasets, LDR leverages target energy information without extra on-policy samples. The method supports unbiased, biased, and variational training regimes and is validated across internal-coordinate flows and Cartesian-coordinate TarFlow implementations, showing substantial improvements in final performance and data efficiency, often by an order of magnitude. The work demonstrates broad applicability and practical impact for efficient equilibrium-ensemble sampling in computational chemistry and physics.

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.
Paper Structure (71 sections, 4 theorems, 44 equations, 7 figures, 13 tables, 1 algorithm)

This paper contains 71 sections, 4 theorems, 44 equations, 7 figures, 13 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. Experiments
  5. Experimental Setup
  6. Architecture
  7. Target densities
  8. Metrics
  9. Training on Unbiased Datasets
  10. Results
  11. Training on Biased Datasets
  12. Dataset
  13. Results
  14. Variational Training without Target Data
  15. Results
  16. ...and 56 more sections

Key Result

Proposition 1.1

Let $r_X$ be a reference distribution with full support on $X$, i.e., $r_X(x)>0$ for all $x\in X$. Let $\tilde{p}_X(x)>0$ be an unnormalized target density and let $q_X^\theta$ be a normalized model density with $q_X^\theta(x)>0$ for all $x$. Then $\mathcal{L}_{\mathrm{LD}}^{\theta\,(p)}(r_X)\ge 0$, In particular, $\mathcal{L}_{\mathrm{LD}}^{\theta\,(p)}(r_X)$ is a divergence whose unique minimum

Figures (7)

  • Figure 1: (a) The log-dispersion objective minimizes the dispersion of $f^\theta(x)=-\log q_X^\theta(x) -(-\log \tilde{p}_X(x))$ around its mean, regularizing the shape of the proposal $q_X^\theta$ over the support of a reference distribution $r_X$. (b) When training off-policy, e.g., with a fixed dataset, log-dispersion alone is not a divergence because it does not constrain the proposal outside the support of $r_X$. We therefore use log-dispersion as a shape regularizer on top of data-based divergences to ensure global normalization. Since the reference distribution $r_X$ is flexible, LDR can be applied to target samples $x_i\sim p_X$ as well as biased samples $x_i^{\text{biased}}\sim p_X^{\text{biased}}$.
  • Figure 2: Ramachandran plot of the biased training dataset for alanine dipeptide. The four starting configurations of the trajectories are labeled in red. Due to the short length of the trajectories, the high-energy metastable states on the right side are oversampled compared to the ground truth.
  • Figure 3: Ramachandran plots for alanine dipeptide, using randomly chosen 1.0e4 samples (left), 2.5e5 (middle), and 1.0e7 samples (right) from the ground truth dataset.
  • Figure 4: Final ESS and NLL as a function of the loss weight $\lambda_\text{data}$ for unbiased training on alanine dipeptide using 1.0e6 samples.
  • Figure 5: Ramachandran obtained for alanine dipeptide when training with only LDR-L1, without an additional data-based divergence.
  • ...and 2 more figures

Theorems & Definitions (8)

  • Proposition 1.1: Log-dispersion regularization as a divergence
  • proof
  • Proposition 1.2: LD without full support is not a divergence
  • proof
  • Proposition 1.3: Consistency of log-dispersion regularization
  • proof
  • Lemma 3.1: Center-of-mass augmentation
  • proof