Value Gradient Sampler: Sampling as Sequential Decision Making

2025-02-18

AI Summary

This work addresses the problem of efficient sampling from a given unnormalized energy function (i.e., an unnormalized density). The method models sampling as a discrete-time sequential decision-making process, where particles—initialized stochastically—are evolved via learned drift and diffusion terms; the drift policy is optimized jointly to minimize an upper bound on the KL divergence between the particle distribution and the target density. Crucially, this is the first framework to rigorously integrate optimal control theory and value-function gradients into sampling, establishing a formal theoretical connection between sampling and reinforcement learning. The approach unifies value-function-based dynamic programming, energy-based modeling, and particle dynamics, enabling end-to-end differentiable training. Empirically, it achieves state-of-the-art performance across multiple standard sampling benchmarks. Moreover, it successfully replaces MCMC in industrial anomaly detection pipelines for training energy-based models, yielding significant improvements in both accuracy and sampling efficiency.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Probabilistic Circuits and Graphical Models
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• 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

Abstract

We propose the Value Gradient Sampler (VGS), a trainable sampler based on the interpretation of sampling as discrete-time sequential decision-making. VGS generates samples from a given unnormalized density (i.e., energy) by drifting and diffusing randomly initialized particles. In VGS, finding the optimal drift is equivalent to solving an optimal control problem where the cost is the upper bound of the KL divergence between the target density and the samples. We employ value-based dynamic programming to solve this optimal control problem, which gives the gradient of the value function as the optimal drift vector. The connection to sequential decision making allows VGS to leverage extensively studied techniques in reinforcement learning, making VGS a fast, adaptive, and accurate sampler that achieves competitive results in various sampling benchmarks. Furthermore, VGS can replace MCMC in contrastive divergence training of energy-based models. We demonstrate the effectiveness of VGS in training accurate energy-based models in industrial anomaly detection applications.

Problem

• Trainable sampler for sequential decision-making
• Optimal drift via value-based dynamic programming
• Enhances energy-based models in anomaly detection

Innovation

• Trainable sampler using decision-making
• Solves optimal control for sampling
• Replaces MCMC in model training