Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

One-Step Sampler for Boltzmann Distributions via Drifting

Wenhan Cao, Keyu Yan, Lin Zhao

Abstract

We present a drifting-based framework for amortized sampling of Boltzmann distributions defined by energy functions. The method trains a one-step neural generator by projecting samples along a Gaussian-smoothed score field from the current model distribution toward the target Boltzmann distribution. For targets specified only up to an unknown normalization constant, we derive a practical target-side drift from a smoothed energy and use two estimators: a local importance-sampling mean-shift estimator and a second-order curvature-corrected approximation. Combined with a mini-batch Gaussian mean-shift estimate of the sampler-side smoothed score, this yields a simple stop-gradient objective for stable one-step training. On a four-mode Gaussian-mixture Boltzmann target, our sampler achieves mean error $0.0754$, covariance error $0.0425$, and RBF MMD $0.0020$. Additional double-well and banana targets show that the same formulation also handles nonconvex and curved low-energy geometries. Overall, the results support drifting as an effective way to amortize iterative sampling from Boltzmann distributions into a single forward pass at test time.

One-Step Sampler for Boltzmann Distributions via Drifting

Abstract

We present a drifting-based framework for amortized sampling of Boltzmann distributions defined by energy functions. The method trains a one-step neural generator by projecting samples along a Gaussian-smoothed score field from the current model distribution toward the target Boltzmann distribution. For targets specified only up to an unknown normalization constant, we derive a practical target-side drift from a smoothed energy and use two estimators: a local importance-sampling mean-shift estimator and a second-order curvature-corrected approximation. Combined with a mini-batch Gaussian mean-shift estimate of the sampler-side smoothed score, this yields a simple stop-gradient objective for stable one-step training. On a four-mode Gaussian-mixture Boltzmann target, our sampler achieves mean error , covariance error , and RBF MMD . Additional double-well and banana targets show that the same formulation also handles nonconvex and curved low-energy geometries. Overall, the results support drifting as an effective way to amortize iterative sampling from Boltzmann distributions into a single forward pass at test time.
Paper Structure (20 sections, 41 equations, 2 figures, 1 table)

This paper contains 20 sections, 41 equations, 2 figures, 1 table.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Preliminaries
  4. Drifting
  5. Smoothed-score drifting
  6. Problem setup: amortized sampling for Boltzmann distributions
  7. Method
  8. Target-side drift for Boltzmann distributions
  9. Interpreting the smoothed energy gradient
  10. Approximating the target-side drift
  11. Monte Carlo approximation.
  12. Second-order approximation.
  13. Estimating the sampler-side smoothed score
  14. Training objective
  15. Mini-batch implementation
  16. ...and 5 more sections

Figures (2)

  • Figure 1: Qualitative behavior of Gaussian kernel drifting on the Gaussian-mixture Boltzmann target. The learned one-step sampler captures all four modes at the correct locations and with approximately the correct spread.
  • Figure 2: Additional target examples for Gaussian kernel drifting. On the left, the learned sampler captures the two-well structure of the double-well energy. On the right, it follows the curved low-energy manifold of the banana-shaped target.