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

2024-09-11

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.

Machine Learning

What problem does this paper attempt to address?

The main problems this paper attempts to address are:

1. **Sampling Bias**: Existing diffusion models (such as DDPM) have inherent errors when generating samples, which are caused by the mismatch between the model and the true distribution.
2. **High Computational Cost**: To obtain high-quality samples, existing methods usually require a large number of function evaluations (NFEs), leading to high computational costs.

Specifically, the paper points out the following two main challenges with current methods:
- **Sampling Bias**: Due to the imperfection of the model, the generated samples have biases that affect the accuracy of the estimates.
- **High Computational Cost**: To obtain high-quality samples, existing methods (such as DDPM) require hundreds of function evaluations (NFEs).

Although existing solutions (such as importance sampling and distillation techniques) can address these issues separately, they are often incompatible. For example, most distillation models lack the necessary density information and cannot be directly used for importance sampling.

To address these issues, the paper proposes a new sampling method that combines Consistency Models (CMs) with importance sampling, thereby maintaining unbiased estimation while reducing computational costs. This method was evaluated on synthetic energy functions and equivariant multi-body particle systems, and the results show that it can generate unbiased samples with only 6-25 NFEs, whereas traditional DDPM methods require about 100 NFEs.