Value Gradient Sampler: Sampling as Sequential Decision Making

TL;DR

The paper introduces Value Gradient Sampler (VGS), a discrete-time, RL-inspired sampler that treats sampling as a sequential decision problem. By solving a value-function-based optimal control problem, VGS computes the drift at each step as the gradient of the next-step value function, enabling fast and accurate sampling from unnormalized densities and offering a drop-in replacement for MCMC in EBM training. Theoretical results link the optimal value function to a diffused auxiliary distribution and establish invariances for symmetric n-body systems; empirically, VGS achieves competitive or superior performance on synthetic distributions, n-body benchmarks, and energy-based anomaly detection tasks, often with fewer time steps than SDE-based methods. The approach provides a principled RL-based framework for sampling that can leverage established RL techniques and symmetry considerations to improve efficiency and energy estimation quality in EBMs.

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.

Figures (4)

• Figure 1: VGS trained on a mixture of 9 Gaussians with $T=10$. The value functions and corresponding samples at each time step are shown. Samples drift along the gradient of the next-step value function. The experimental setup and results are explained in \ref{['sec:synthetic_exp']}.
• Figure 2: Sinkhorn distance $\mathcal{W}_{\gamma}^{2}$ to the number of time steps $T$ on GMM. Performance of VGS shows robustness to the number of time steps $T$ and outperforms DDS on every $T<100$. Increasing the noise magnitude $s_t$ in VGS improves performance in the small $T$ regime but makes training divergent when $T$ is large.
• Figure 3: Energy (Left) and Interatomic Distance (Right) histograms of DW-4 samples from test data, VGS, and VGS with a non-invariant network. VGS generates accurate samples by leveraging the symmetry of the system.
• Figure 4: EBM training on 2D 8 Gaussians. The red shade depicts the energy, and the dots are the samples. VGS produces an accurate energy estimate, while the short-run MCMC does not.

Theorems & Definitions (8)

• Theorem 3.1: Optimal Value Function
• proof
• Theorem 3.2: Optimal Auxiliary Distribution
• proof
• Theorem 4.1: Invariance of the Value Function
• proof
• proof
• proof