Designing diffusion models for tractable likelihood estimation

Estimating likelihoods with diffusion models is typically done by expensive approximations of the divergence of a neural network, as exact computation is intractable in general. Learning the time score ${\partial }_t\log p$ is a recent idea that may enable more viable and exact computation, but is still a challenging task as previous work has been limited in scale, has suffered with sampling degradation and has mostly focused on the energy setting, which by itself already suffers from training difficulties. In this work, we tackle these issues by proposing a unifying framework to design diffusion models for effective sampling and likelihood estimation. First, we propose a regularization strategy to incorporate the time score without negatively affecting sampling. Then, we introduce a set of design coefficients and learnable terms that generalizes and connects important cases, such as the energy and mean flow formulations. Finally, specific parametrization and learning strategies are proposed in order to avoid singularities and improve training. The resulting algorithms are tested and studied on synthetic examples and on real images for ImageNet on both pixel and latent space.