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

TL;DR

This work investigates diffusion-based annealed Monte Carlo Boltzmann Generators (BGs) for sampling from complex Boltzmann targets. It analyzes two regimes: an idealized setting with perfectly learned diffusion models and a realistic setting where DM components are learned from data. Key findings show that, with perfect knowledge, diffusion-induced density paths and second-order or deterministic transitions substantially improve DM-aMC BGs over first-order, but when logs/densities are learned, BGs fail due to mode blindness in log-density estimation rather than score quality. The results motivate developing deterministic transport maps and robust density-learning objectives to realize reliable diffusion-based BGs, highlighting mode-switching as a central challenge and outlining concrete paths for future work in adaptive training and mode-aware density modeling. Overall, the study delineates both the potential benefits and the crucial bottlenecks of DM-based BGs, guiding future refinements toward practical, scalable Boltzmann sampling in multimodal, high-dimensional settings.

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.

Figures (63)

• Figure 1: Sampling results for classic annealed samplers with diffusion (blue) and tempering (red) density paths, when targeting TwoModes(Top) and ManyModes(Bottom) distributions in idealized setting \ref{['item:setting_perfect']}. For tempering paths, we display the best-performing result among all values of $K$. For diffusion paths, we display the results for all values of $K$ : the darker the bar, the higher $K$. In particular, these configurations do not share the same computational budget. Each result is averaged over 8 runs with 8,192 samples per run. We observe that for almost all target settings and samplers, there exists a number of annealing levels $K$ for which the noising density path outperforms the best tempering baseline.
• Figure 2: Different diffusion-based swapping mechanisms for Replica Exchange.(Left) standard swap scheme, see \ref{['sec:monte_carlo']}, where samples are exchanged directly across noise levels without guidance, potentially moving into low-probability regions. (Middle) DM-based swaps using forward and backward Markov kernels, as proposed by zhang2025acceleratedparalleltemperingneural (coined Diff-APT), see \ref{['subsec:review_methods']}. (Right) DM-based swaps using forward and backward transport maps under the deterministic framework introduced in \ref{['subsec:determinist_general']}. In each panel, black lines denote noise levels; green dots mark the original samples; blue and red paths indicate forward (low to high noise) and backward (high to low noise) trajectories, respectively; and the swapped samples are shown as the resulting blue and green dots. The DM-based swaps better preserve high-probability regions during exchange, enabling theoretically more effective sampling.
• Figure 3: DM-based aMC-BG results with annealed samplers using different mechanisms, when targeting TwoModes distribution in idealized setting \ref{['item:setting_perfect']} : (First row) Low distance, high dimension, (Second row) High distance, low dimension, (Third row) Middle distance, middle dimension, (Fourth row) Average running time of each sampler over the three TwoModes variants. Each group of bars with the same color corresponds to a specific aMC method. Within each group, red bars refer to methods that do not use DM-based transition kernels (standard baseline); blue and green refer to methods that exploit 1st-order and 2nd-order stochastic kernels (prior work); pink and yellow correspond to variants with deterministic maps, where the log-determinant term is respectively computed from the ground truth diagonal Hessian or via the Hutchinson trick. For all settings, we display the results for all values of $K$ : the darker the bar, the higher $K$ (same range as \ref{['fig:temp_vs_noising']}). In particular, configurations within each bar group do not share the same computational budget. Each result is averaged over 8 runs with 8,192 samples per run. We observe that using 1st-order stochastic kernels (in blue) does not always lead to better performance than the baseline (in red), while 2nd-order stochastic kernels (in green) provide consistent improvements. Moreover, using deterministic transitions combined with the Hutchinson trick (in yellow) brings better performance than 1st-order stochastic kernels (both only rely on the use of the scores), while having access to the Hessians (in pink) performs comparably to the 2nd-order stochastic methods. In the case of AIS/SMC sampler, using deterministic transitions incur only a little computational overhead compared to the stochastic case, but this cost is stringer for RE sampler. Similar results are given in \ref{['app:bonus_exp']} for ManyModes targets.
• Figure 4: Realistic results of DM-based aMC-BG, when targeting TwoModes distribution with intermediate difficulty (middle mode distance, middle dimension) in setting \ref{['item:setting_practical']} : (From top to bottom) the DM is trained via TSM+DSM, tSM+DSM, aLFPE+DSM or RNE+DSM objective with identical computational budget. Each group of bars with the same color corresponds to a specific aMC method, except for the last two groups on the right which shows the baseline obtained by directly simulating the reverse SDE \ref{['eq:sde-denoising']} and reverse ODE \ref{['eq:pf_ode']}. Bar colors indicate the type of density path used for the sampling method, and are consistent with those displayed in \ref{['fig:diff_two_modes']}: purple for the model-free tempering path (same baseline for all models), red for the diffusion density path solely exploiting the log-densities, blue, respectively yellow, for the diffusion density path additionally exploiting DM first-order stochastic kernels, respectively DM first-order deterministic kernels. On the other hand, hatching denotes the nature of log densities used in aMC (exact as in \ref{['fig:diff_two_modes']}, learned with hardcoded EBM or learned with pinned EBM). For each method and density path, the number of levels $K$ was optimized individually so as to display the best expected result from the considered approach if $K$ was carefully tuned. The results show a clear performance gap between learned and ideal paths. Notably, DM-BGs generally underperform compared to directly to the DM alone (i.e. simulating the reverse SDE/ODE), and rarely surpass the classic tempering methods. Similar results are given in \ref{['app:bonus_exp']} for other TwoModes and ManyModes targets, and the other DM training objectives.
• Figure 5: Exact density paths bridging $\pi^{\text{base}}$ (last time index) to 1D Gaussian mixtures (first time index).(Left): the target is an instance of TwoModes defined as $(3/4) \mathrm{N}(-4, 0.5^2) + (1/4) \mathrm{N}(+4, 1)$, (Right) the target is the 1D instance of ManyModesNoble2024learned with 32 modes, (First and third columns) diffusion density path, (Second and fourth columns) tempering density path. Given a discretization of 128 timesteps in [0,1], we plot each of the induced marginal densities using importance sampling with 1,000,000 particles and a uniform proposal. We observe that the tempering path shows clear mode switching for both of the targets: in the case of the TwoModes target, the strongest mode emerges abruptly, while the weakest modes appear rapidly for the ManyModes target. On the other hand, the mode weights in the diffusion path remain stable over time, making it more favorable for aMC.

Theorems & Definitions (50)

• Proposition 1: IM integrator with Euler scheme
• Proposition 2: Fixed-point approximation of the IM integrator
• Proposition 3: Approximation of the Jacobian log-determinants via power series
• Lemma 1: Power series expansion of the matrix logarithm
• proof
• Corollary 1
• proof
• Lemma 2: SDE exponential integration
• proof
• Lemma 3: ODE exponential integration