Fast Sampling of Diffusion Models via Operator Learning
Diffusion models have found widespread adoption in various areas. However,
sampling from them is slow because it involves emulating a reverse process with
hundreds-to-thousands of network evaluations. Inspired by the 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
models 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, which is a continuous
function in time. DFNO can be applied to any diffusion model and generate
high-quality samples in one model forward call. Our method achieves the
state-of-the-art FID of 4.72 on CIFAR-10 using only one model evaluation.
Authors
Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, Anima Anandkumar