Value Gradient Sampler: Learning Invariant Value Functions for Equivariant Diffusion Sampling

Abstract

We propose the Value Gradient Sampler (VGS), a diffusion sampler parameterized by value functions. VGS generates samples from an unnormalized target density (i.e., energy) by evolving randomly initialized particles along the gradient of the value function. In many sampling problems where the target density exhibits equivariant symmetries, we show that value functions enable a novel approach to leveraging invariant neural networks for sampling, as an invariant value function induces an equivariant gradient flow. The value functions are trained via temporal-difference learning, which supports off-policy training and other established reinforcement learning (RL) techniques. By combining efficient invariant neural networks with advanced RL methods, VGS achieves strong performance in high-dimensional particle systems, including Lennard-Jones systems with up to 55 particles.