Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient and Unbiased Sampling of Boltzmann Distributions via Consistency Models

Fengzhe Zhang, Jiajun He, Laurence I. Midgley, Javier Antorán, José Miguel Hernández-Lobato

TL;DR

This work addresses efficient, unbiased sampling from Boltzmann distributions by marrying Importance Sampling (IS) with Bidirectional Consistency Models (BCMs). The proposed framework defines a proposal and a target distribution along the Pf ODE/SDE trajectory, enabling IS to correct model bias while CM-based traversal keeps NFEs low. The authors extend Consistency Trajectory Models to Bidirectional CTMs with E(3)-equivariance for molecular data and demonstrate unbiased sampling with 6–25 NFEs, achieving ESS comparable to DDPMs that use ~100 NFEs on synthetic and equivariant n-body tasks. The approach offers a practical, density-aware accelerated sampler suitable for Boltzmann generators, with limitations including an observed plateau in ESS gains at higher NFEs and the need for robust hyperparameter tuning in high dimensions.

Abstract

Diffusion models have shown promising potential for advancing Boltzmann Generators. However, two critical challenges persist: (1) inherent errors in samples due to model imperfections, and (2) the requirement of hundreds of functional evaluations (NFEs) to achieve high-quality samples. While existing solutions like importance sampling and distillation address these issues separately, they are often incompatible, as most distillation models lack the necessary density information for importance sampling. This paper introduces a novel sampling method that effectively combines Consistency Models (CMs) with importance sampling. We evaluate our approach on both synthetic energy functions and equivariant n-body particle systems. Our method produces unbiased samples using only 6-25 NFEs while achieving a comparable Effective Sample Size (ESS) to Denoising Diffusion Probabilistic Models (DDPMs) that require approximately 100 NFEs.

Efficient and Unbiased Sampling of Boltzmann Distributions via Consistency Models

TL;DR

This work addresses efficient, unbiased sampling from Boltzmann distributions by marrying Importance Sampling (IS) with Bidirectional Consistency Models (BCMs). The proposed framework defines a proposal and a target distribution along the Pf ODE/SDE trajectory, enabling IS to correct model bias while CM-based traversal keeps NFEs low. The authors extend Consistency Trajectory Models to Bidirectional CTMs with E(3)-equivariance for molecular data and demonstrate unbiased sampling with 6–25 NFEs, achieving ESS comparable to DDPMs that use ~100 NFEs on synthetic and equivariant n-body tasks. The approach offers a practical, density-aware accelerated sampler suitable for Boltzmann generators, with limitations including an observed plateau in ESS gains at higher NFEs and the need for robust hyperparameter tuning in high dimensions.

Abstract

Diffusion models have shown promising potential for advancing Boltzmann Generators. However, two critical challenges persist: (1) inherent errors in samples due to model imperfections, and (2) the requirement of hundreds of functional evaluations (NFEs) to achieve high-quality samples. While existing solutions like importance sampling and distillation address these issues separately, they are often incompatible, as most distillation models lack the necessary density information for importance sampling. This paper introduces a novel sampling method that effectively combines Consistency Models (CMs) with importance sampling. We evaluate our approach on both synthetic energy functions and equivariant n-body particle systems. Our method produces unbiased samples using only 6-25 NFEs while achieving a comparable Effective Sample Size (ESS) to Denoising Diffusion Probabilistic Models (DDPMs) that require approximately 100 NFEs.
Paper Structure (17 sections, 11 equations, 3 figures, 2 tables)

This paper contains 17 sections, 11 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Importance Sampling.
  4. Score-based Diffusion Models.
  5. Consistency Models.
  6. Method
  7. Proposal Distribution.
  8. Target Distribution.
  9. Time Step Optimization.
  10. Model Extensions.
  11. Experiments
  12. Conclusions and Limitations
  13. Baseline: Combining Importance Sampling with Diffusion Models
  14. Efficacy of Alternating ODE-SDE Target Distribution Design
  15. Tuning the Time Steps
  16. ...and 2 more sections

Figures (3)

  • Figure 1: Overview of our proposed sampling method, which combines IS with BCMs.
  • Figure 2: ESS of our proposed algorithm with BCM and BCTM. We also include DDPM as baseline. We estimate ESS with 100,000 samples and report the mean, the first and the third quantile.
  • Figure 3: Visualization of proposal and target distributions for different target distribution designs.