Timezone: »

 
Fast Sampling of Diffusion Models via Operator Learning
Hongkai Zheng · Weili Nie · Arash Vahdat · Kamyar Azizzadenesheli · Anima Anandkumar

@
in NeurIPS 2022 Workshop on Score-Based Methods »
Event URL: https://openreview.net/forum?id=XrhofG6qg7Y »

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.

Author Information

Hongkai Zheng (Caltech)
Weili Nie (NVIDIA)
Arash Vahdat (NVIDIA Research)
Kamyar Azizzadenesheli (Purdue University)
Anima Anandkumar (NVIDIA / Caltech)

More from the Same Authors