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 invariant symmetries, value functions provide a novel approach to leveraging invariant networks for sampling by inducing an equivariant gradient flow, without requiring more complex equivariant networks. The value networks are trained via temporal difference learning, which supports off-policy training and other established reinforcement learning (RL) techniques. By combining advanced RL methods with efficient invariant networks, VGS achieves both the highest sample quality and the fastest sampling speed among our baselines on the 55-particle Lennard-Jones system.

Code Dataset Promise: Yes

Code Dataset Url: https://github.com/swyoon/value-gradient-sampler

Format Confirmation: I agree that I have read and followed the formatting instructions for the camera ready version.

Submission Number: 1719