[Paper Review] Fast Sampling of Diffusion Models via Operator Learning

Hongkai Zheng, Weili Nie+3 more|arXiv (Cornell University)|Nov 24, 2022

Advanced Neuroimaging Techniques and ApplicationsMedicine17 citations

TL;DR

This paper proposes DSNO, a novel neural operator-based method for fast, parallel sampling of diffusion models using a single forward pass. By integrating Fourier-domain temporal convolution layers into a diffusion model backbone, DSNO learns the continuous-time reverse process trajectory, achieving state-of-the-art FID scores of 3.78 on CIFAR-10 and 7.83 on ImageNet-64 with one model evaluation.

ABSTRACT

Diffusion models have found widespread adoption in various areas. However, their sampling process is slow because it requires hundreds to thousands of network evaluations to emulate a continuous process defined by differential equations. In this work, we use neural operators, an efficient method to solve the probability flow differential equations, to accelerate the sampling process of diffusion models. Compared to other fast sampling methods that have a sequential nature, we are the first to propose a parallel decoding method that generates images with only one model forward pass. We propose diffusion model sampling with neural operator (DSNO) that maps the initial condition, i.e., Gaussian distribution, to the continuous-time solution trajectory of the reverse diffusion process. To model the temporal correlations along the trajectory, we introduce temporal convolution layers that are parameterized in the Fourier space into the given diffusion model backbone. We show our method achieves state-of-the-art FID of 3.78 for CIFAR-10 and 7.83 for ImageNet-64 in the one-model-evaluation setting.

Motivation & Objective

• Accelerate diffusion model sampling to enable real-time inference for time-sensitive applications.
• Overcome the sequential bottleneck of existing fast sampling methods that require multiple forward passes.
• Develop a differentiable, continuous-time solution trajectory mapping from noise to data using neural operators.
• Preserve high sample quality while drastically reducing inference time to a single forward pass.
• Integrate temporal modeling into diffusion models using parameter-efficient, Fourier-space temporal convolutions.

Proposed method

• Propose DSNO, a neural operator that maps the initial Gaussian noise to the continuous-time solution trajectory of the reverse diffusion process.
• Introduce temporal convolution layers parameterized in the Fourier domain to model temporal correlations along the trajectory.
• Enable parallel decoding by computing all time-step outputs simultaneously via Fourier-based operations, avoiding sequential computation.
• Integrate the temporal convolution blocks into an existing U-Net backbone, maintaining compatibility with standard diffusion model architectures.
• Train the model using a combination of L1 and perceptual (LPIPS) losses to improve sample fidelity.
• Use quadratic time discretization and multi-resolution supervision to improve generalization and FID performance.

Experimental results

Research questions

• RQ1 Can a neural operator be effectively used to model the continuous-time reverse process of diffusion models for fast sampling?
• RQ2 Can temporal convolution layers in Fourier space capture long-range temporal dependencies in the diffusion trajectory efficiently?
• RQ3 Does parallel decoding via Fourier-based operations enable one-step generation without sacrificing sample quality?
• RQ4 Can the proposed method achieve state-of-the-art FID scores with only a single forward pass?
• RQ5 How does the choice of loss function (e.g., L1 vs. LPIPS) affect the final sample quality and FID score?

Key findings

• DSNO achieves a state-of-the-art FID score of 3.78 on CIFAR-10 under the one-model-evaluation setting.
• DSNO achieves a state-of-the-art FID score of 7.83 on ImageNet-64 with only one model evaluation.
• The use of VGG-based LPIPS loss leads to a measurable improvement in FID compared to standard L1 loss.
• Increasing temporal resolution from 4 to 8 improves FID slightly, but the marginal gain suggests 4 is sufficient for efficiency.
• The proposed temporal convolution blocks add only 10% to the model size while enabling parallel decoding and high performance.
• The method enables true parallel decoding, generating all time-step outputs in a single forward pass, unlike sequential solvers.