Neural Diffusion Models

TL;DR

Neural Diffusion Models extend traditional diffusion models by enabling non-linear, time-dependent transformations, leading to improved likelihoods and sample quality in image generation tasks.

Contribution

The paper introduces Neural Diffusion Models, allowing non-linear transformations and providing a variational training method with a continuous-time formulation.

Findings

NDMs outperform conventional diffusion models on likelihood metrics.

NDMs generate higher quality images on benchmark datasets.

The continuous formulation enables efficient inference with standard solvers.

Abstract

Diffusion models have shown remarkable performance on many generative tasks. Despite recent success, most diffusion models are restricted in that they only allow linear transformation of the data distribution. In contrast, broader family of transformations can potentially help train generative distributions more efficiently, simplifying the reverse process and closing the gap between the true negative log-likelihood and the variational approximation. In this paper, we present Neural Diffusion Models (NDMs), a generalization of conventional diffusion models that enables defining and learning time-dependent non-linear transformations of data. We show how to optimise NDMs using a variational bound in a simulation-free setting. Moreover, we derive a time-continuous formulation of NDMs, which allows fast and reliable inference using off-the-shelf numerical ODE and SDE solvers. Finally, we demonstrate the utility of NDMs with learnable transformations through experiments on standard image generation benchmarks, including CIFAR-10, downsampled versions of ImageNet and CelebA-HQ. NDMs outperform conventional diffusion models in terms of likelihood and produce high-quality samples.