Fast Sampling of Diffusion Models via Operator Learning

Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, Anima Anandkumar

Keywords: sampling, diffusion model, score-based generative model, neural operator

TL;DR: We propose diffusion Fourier neural operator (DFNO) with temporal convolution in Fourier space for fast sampling of score-based models.

Abstract: Diffusion models have found widespread adoption in various areas. However, sampling from them is still slow because it involves emulating a reverse stochastic process with hundreds-to-thousands of neural network evaluations. Inspired by the recent success of neural operators in accelerating differential equations solving, we approach this problem by solving the underlying neural differential equation from an operator learning perspective. We examine probability flow ODE trajectories in diffusion model and observe a compact energy spectrum that can be learned efficiently in Fourier space. With this insight, we propose diffusion Fourier neural operator (DFNO) with temporal convolution in Fourier space to parameterize the operator that maps initial condition to the solution trajectory. DFNO can apply to any diffusion models and generate high-quality samples in one step. Our method achieves the state-of-the-art clean FID of 5.9 (legacy FID 4.72) on CIFAR-10 using one network evaluation.

Student Paper: Yes