One-Step Sampler for Boltzmann Distributions via Drifting

2026-03-18

AI Summary

This work addresses the high computational cost of sampling from Boltzmann distributions during inference by proposing a drift-based amortized sampling framework. The method trains a single-step neural generator that projects current samples onto the target distribution in one step along a Gaussian-smoothed score field, eliminating the need for iterative refinement and accommodating complex energy landscapes with unknown normalizing constants, non-convexity, or curved low-energy structures. Innovatively, it combines a mean-shift estimator from local importance sampling with second-order curvature correction to construct a target-side drift term, while employing a stop-gradient strategy to ensure stable training. Experiments demonstrate superior performance on multimodal Gaussian mixtures, double-well, and banana-shaped distributions—achieving RBF MMD as low as 0.0020—and validate the approach’s efficacy in efficiently adapting to intricate energy landscapes.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

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.

Problem

• Boltzmann distributions
• amortized sampling
• energy-based models
• normalization constant
• non-iterative sampling

Innovation

• drifting
• one-step sampling
• Boltzmann distribution
• score matching
• amortized inference