English

One-Step Sampler for Boltzmann Distributions via Drifting

Machine Learning 2026-03-19 v1

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.07540.0754, covariance error 0.04250.0425, and RBF MMD 0.00200.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.

Keywords

Cite

@article{arxiv.2603.17579,
  title  = {One-Step Sampler for Boltzmann Distributions via Drifting},
  author = {Wenhan Cao and Keyu Yan and Lin Zhao},
  journal= {arXiv preprint arXiv:2603.17579},
  year   = {2026}
}
R2 v1 2026-07-01T11:25:55.599Z