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

TL;DR

The paper introduces the Value Gradient Sampler (VGS), a diffusion sampling method using value functions that leverages invariant networks and reinforcement learning techniques to improve sample quality and speed.

Contribution

VGS is a novel diffusion sampler that employs invariant value functions and RL training to efficiently generate samples with invariant properties.

Findings

VGS achieves the highest sample quality among baselines.

VGS provides the fastest sampling speed.

VGS effectively leverages invariant networks for 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 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.