One-Step Sampler for Boltzmann Distributions via Drifting

TL;DR

This paper introduces a drifting-based framework for efficiently sampling Boltzmann distributions using a neural generator trained with a novel one-step approach, enabling fast and stable approximate sampling.

Contribution

It proposes a new one-step neural sampling method for Boltzmann distributions that handles unknown normalization constants and complex geometries.

Findings

Achieves low mean and covariance errors on Gaussian-mixture targets.

Effectively handles nonconvex and curved low-energy landscapes.

Enables single-pass sampling at test time with stable training.

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.