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.

@inproceedings{TongICLR25,
  author    = {Tong, Vinh and Liu, Anji and Hoang, Trung-Dung and Van den Broeck, Guy and Niepert, Mathias},
  title     = {Learning to Discretize Denoising Diffusion ODEs}, 
  booktitle = {Proceedings of the 13th International Conference on Learning Representations (ICLR)},
  month     = 4,
  year      = {2025},
  url       = "https://starai.cs.ucla.edu/papers/TongICLR25.pdf",
  annotation = "(Oral full presentation, acceptance rate 1.8\%)",
  keywords  = {conference,selective}
}