Improved off-policy training of diffusion samplers

Mila - Quebec AI Institute
1Université de Montréal 2Jagiellonian University 3KAIST
4Ciela Institute 5Center for Computational Astrophysics, Flatiron Institute
6Dreamfold 7CIFAR AI Chair
Teaser image.

Abstract

We study the problem of training diffusion models to sample from a distribution with a given unnormalized density or energy function. We benchmark several diffusion-structured inference methods, including simulation-based variational approaches and off-policy methods (continuous generative flow networks). Our results shed light on the relative advantages of existing algorithms while bringing into question some claims from past work. We also propose a novel exploration strategy for off-policy methods, based on local search in the target space with the use of a replay buffer, and show that it improves the quality of samples on a variety of target distributions. Our code for the sampling methods and benchmarks studied is made public at (link) as a base for future work on diffusion models for amortized inference.



BibTeX

 
      @inproceedings{
        sendera2024improved,
        title={Improved off-policy training of diffusion samplers},
        author={Marcin Sendera and Minsu Kim and Sarthak Mittal and Pablo Lemos and Luca Scimeca and Jarrid Rector-Brooks and Alexandre Adam and Yoshua Bengio and Nikolay Malkin},
        booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
        year={2024},
        url={https://openreview.net/forum?id=vieIamY2Gi}
      }