Research

Paper

TESTING March 18, 2026

One-Step Sampler for Boltzmann Distributions via Drifting

Authors

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.

Metadata

arXiv ID: 2603.17579
Provider: ARXIV
Primary Category: cs.LG
Published: 2026-03-18
Fetched: 2026-03-19 06:01

Related papers

Raw Data (Debug)
{
  "raw_xml": "<entry>\n    <id>http://arxiv.org/abs/2603.17579v1</id>\n    <title>One-Step Sampler for Boltzmann Distributions via Drifting</title>\n    <updated>2026-03-18T10:35:16Z</updated>\n    <link href='https://arxiv.org/abs/2603.17579v1' rel='alternate' type='text/html'/>\n    <link href='https://arxiv.org/pdf/2603.17579v1' rel='related' title='pdf' type='application/pdf'/>\n    <summary>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.</summary>\n    <category scheme='http://arxiv.org/schemas/atom' term='cs.LG'/>\n    <published>2026-03-18T10:35:16Z</published>\n    <arxiv:primary_category term='cs.LG'/>\n    <author>\n      <name>Wenhan Cao</name>\n    </author>\n    <author>\n      <name>Keyu Yan</name>\n    </author>\n    <author>\n      <name>Lin Zhao</name>\n    </author>\n  </entry>"
}