Training Neural Samplers with Reverse Diffusive KL Divergence

TL;DR

This paper introduces a novel training method for neural samplers that minimizes reverse diffusive KL divergence along diffusion trajectories, enabling effective approximation of multi-modal distributions in a single step.

Contribution

The authors propose the reverse diffusive KL divergence objective, improving neural sampler training for multi-modal distributions compared to traditional reverse KL methods.

Findings

Enhanced sampling of multi-modal distributions

Effective one-step generation of samples

Improved performance on Boltzmann distributions

Abstract

Training generative models to sample from unnormalized density functions is an important and challenging task in machine learning. Traditional training methods often rely on the reverse Kullback-Leibler (KL) divergence due to its tractability. However, the mode-seeking behavior of reverse KL hinders effective approximation of multi-modal target distributions. To address this, we propose to minimize the reverse KL along diffusion trajectories of both model and target densities. We refer to this objective as the reverse diffusive KL divergence, which allows the model to capture multiple modes. Leveraging this objective, we train neural samplers that can efficiently generate samples from the target distribution in one step. We demonstrate that our method enhances sampling performance across various Boltzmann distributions, including both synthetic multi-modal densities and n-body particle systems.