Diffusion-based Annealed Boltzmann Generators : benefits, pitfalls and hopes

TL;DR

This paper evaluates diffusion-based annealed Boltzmann Generators for sampling from complex distributions, highlighting their benefits, limitations, and potential improvements through empirical analysis and novel deterministic integration methods.

Contribution

It introduces a deterministic aMC integration method using transport maps from diffusion models and provides a detailed empirical analysis of DM-based aMC-BGs in multi-modal scenarios.

Findings

Second-order denoising kernels improve performance with covariance info.

Deterministic aMC with transport maps outperforms stochastic methods at higher costs.

Inaccurate DM log-density estimation is a key bottleneck in learned models.

Abstract

Sampling configurations at thermodynamic equilibrium is a central challenge in statistical physics. Boltzmann Generators (BGs) tackle it by combining a generative model with a Monte Carlo (MC) correction step to obtain asymptotically unbiased samples from an unnormalized target. Most current BGs use classic MC mechanisms such as importance sampling, which both require tractable likelihoods from the backbone model and scale poorly in high-dimensional, multi-modal targets. We study BGs built on annealed Monte Carlo (aMC), which is designed to overcome these limitations by bridging a simple reference to the target through a sequence of intermediate densities. Diffusion models (DMs) are powerful generative models and have already been incorporated into aMC-based recalibration schemes via the diffusion-induced density path, making them appealing backbones for aMC-BGs. We provide an empirical meta-analysis of DM-based aMC-BGs on controlled multi-modal Gaussian mixtures (varying mode separation, number of modes, and dimension), explicitly disentangling inference effects from learning effects by comparing (i) a perfectly learned DM and (ii) a DM trained from data. Even with a perfect DM, standard integrations using only first-order stochastic denoising kernels fail systematically, whereas second-order denoising kernels can substantially improve performance when covariance information is available. We further propose a deterministic aMC integration based on first-order transport maps derived from DMs, which outperforms the stochastic first-order variant at higher computational cost. Finally, in the learned-DM setting, all DM-aMC variants struggle to produce accurate BGs; we trace the main bottleneck to inaccurate DM log-density estimation.