Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient and Unbiased Sampling from Boltzmann Distributions via Variance-Tuned Diffusion Models

Fengzhe Zhang, Laurence I. Midgley, José Miguel Hernández-Lobato

TL;DR

Variance-Tuned Diffusion Importance Sampling (VT-DIS), a lightweight post-training method that adapts the per-step noise covariance of a pretrained SBDM by minimizing the $\alpha$-divergence between its forward diffusion and reverse denoising trajectories, is introduced.

Abstract

Score-based diffusion models (SBDMs) are powerful amortized samplers for Boltzmann distributions; however, imperfect score estimates bias downstream Monte Carlo estimates. Classical importance sampling (IS) can correct this bias, but computing exact likelihoods requires solving the probability-flow ordinary differential equation (PF-ODE), a procedure that is prohibitively costly and scales poorly with dimensionality. We introduce Variance-Tuned Diffusion Importance Sampling (VT-DIS), a lightweight post-training method that adapts the per-step noise covariance of a pretrained SBDM by minimizing the $α$-divergence ($α=2$) between its forward diffusion and reverse denoising trajectories. VT-DIS assigns a single trajectory-wise importance weight to the joint forward-reverse process, yielding unbiased expectation estimates at test time with negligible overhead compared to standard sampling. On the DW-4, LJ-13, and alanine-dipeptide benchmarks, VT-DIS achieves effective sample sizes of approximately 80 %, 35 %, and 3.5 %, respectively, while using only a fraction of the computational budget required by vanilla diffusion + IS or PF-ODE-based IS.

Efficient and Unbiased Sampling from Boltzmann Distributions via Variance-Tuned Diffusion Models

TL;DR

Variance-Tuned Diffusion Importance Sampling (VT-DIS), a lightweight post-training method that adapts the per-step noise covariance of a pretrained SBDM by minimizing the -divergence between its forward diffusion and reverse denoising trajectories, is introduced.

Abstract

Score-based diffusion models (SBDMs) are powerful amortized samplers for Boltzmann distributions; however, imperfect score estimates bias downstream Monte Carlo estimates. Classical importance sampling (IS) can correct this bias, but computing exact likelihoods requires solving the probability-flow ordinary differential equation (PF-ODE), a procedure that is prohibitively costly and scales poorly with dimensionality. We introduce Variance-Tuned Diffusion Importance Sampling (VT-DIS), a lightweight post-training method that adapts the per-step noise covariance of a pretrained SBDM by minimizing the -divergence () between its forward diffusion and reverse denoising trajectories. VT-DIS assigns a single trajectory-wise importance weight to the joint forward-reverse process, yielding unbiased expectation estimates at test time with negligible overhead compared to standard sampling. On the DW-4, LJ-13, and alanine-dipeptide benchmarks, VT-DIS achieves effective sample sizes of approximately 80 %, 35 %, and 3.5 %, respectively, while using only a fraction of the computational budget required by vanilla diffusion + IS or PF-ODE-based IS.

Paper Structure

This paper contains 47 sections, 2 theorems, 43 equations, 7 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Background
  4. Score-Based Diffusion Models
  5. E(3)–Equivariant Score-Based Diffusion Models
  6. Diffusion Models with Importance Sampling
  7. Method
  8. Overview
  9. Covariance Parameterization
  10. Isotropic Covariance
  11. Structured covariances
  12. Sample-dependent models
  13. Extension to E(3)-Equivariant Diffusion
  14. Molecular distributions.
  15. Objective Function
  16. ...and 32 more sections

Key Result

Lemma 1

Let the reverse transition be where $B\in\mathbb{R}^{M\times M}$ is positive definite. Then the reverse kernel is $\mathrm{E}(3)$-equivariant.

Figures (7)

  • Figure 1: GMM-2: ESS versus NFEs for dimensions $d=50$ and $d=100$, comparing VT-DIS with the DDPM + IS baseline. ESS estimates are computed using $10^5$ samples.
  • Figure 2: ESS versus NFEs. VT-DIS consistently outperforms the standard DDPM + IS baseline across all tasks under the same computational budget. ESS values are estimated using $10^5$ samples.
  • Figure 3: ELBO/EUBO vs. diffusion steps. ELBO (higher is better) and EUBO (lower is better) for the DDPM baseline and VT-DIS (isotropic, diagonal and full covariances) as a function of the number of sampling steps. Error bars denote $\pm$ one standard deviation across 3 repeated MC evaluations.
  • Figure 4: Energy distributions. Unweighted samples generated by VT-DIS with isotropic covariance (orange) are biased relative to ground-truth samples (blue), but importance-weighted samples (red) effectively correct the discrepancy. Histograms are plotted using $5\times10^5$ samples.
  • Figure 5: Ramachandran plots from left to right: ground-truth samples generated by MD, VT-DIS (isotropic) samples before and after reweighting. All plots are generated using $10^6$ samples.
  • ...and 2 more figures

Theorems & Definitions (4)

  • Lemma 1
  • proof
  • Lemma 2: Permutation equivariance
  • proof