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

Himchan Hwang, Hyeokju Jeong, Dong Kyu Shin, Che-Sang Park, Sehee Kweon, Sangwoong Yoon, Frank C. Park

TL;DR: Value Gradient Sampler (VGS) is a diffusion sampler parametrized by value functions and trained via temporal difference learning. VGS is particularly effective and efficient in sampling from densities with invariance symmetries.

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.

Submission Number: 1719