Learning to Discretize Denoising Diffusion ODEs

TL;DR

This paper introduces LD3, a lightweight framework that learns optimal time discretization to improve sampling efficiency in diffusion probabilistic models, reducing computational costs while maintaining high-quality generation.

Contribution

LD3 is a novel method that learns the optimal discretization for sampling in diffusion models without retraining the entire network.

Findings

LD3 improves sampling efficiency with less computational overhead.

Achieves state-of-the-art FID scores with fewer neural function evaluations.

Demonstrates effectiveness across multiple pre-trained models and sampling settings.

Abstract

Diffusion Probabilistic Models (DPMs) are generative models showing competitive performance in various domains, including image synthesis and 3D point cloud generation. Sampling from pre-trained DPMs involves multiple neural function evaluations (NFEs) to transform Gaussian noise samples into images, resulting in higher computational costs compared to single-step generative models such as GANs or VAEs. Therefore, reducing the number of NFEs while preserving generation quality is crucial. To address this, we propose LD3, a lightweight framework designed to learn the optimal time discretization for sampling. LD3 can be combined with various samplers and consistently improves generation quality without having to retrain resource-intensive neural networks. We demonstrate analytically and empirically that LD3 improves sampling efficiency with much less computational overhead. We evaluate our method with extensive experiments on 7 pre-trained models, covering unconditional and conditional sampling in both pixel-space and latent-space DPMs. We achieve FIDs of 2.38 (10 NFE), and 2.27 (10 NFE) on unconditional CIFAR10 and AFHQv2 in 5-10 minutes of training. LD3 offers an efficient approach to sampling from pre-trained diffusion models. Code is available at https://github.com/vinhsuhi/LD3.