PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast Sampling of Diffusion Models via Operator Learning

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

arXiv:2211.13449v3cs.LGcs.CV

TL;DR

Diffusion models sample slowly because emulating their continuous-time dynamics requires many network evaluations. DSNO learns a neural-operator map from Gaussian initial conditions to the reverse-diffusion trajectory, using Fourier temporal convolutions for parallel decoding. It generates images with one model evaluation and reports FID scores of 3.78 on CIFAR-10 and 7.83 on ImageNet-64.

  • Problem

    Diffusion-model sampling is slow because emulating the continuous reverse-diffusion process requires hundreds to thousands of network evaluations.

  • Method

    DSNO uses a neural operator with Fourier-parameterized temporal convolution layers to map Gaussian initial conditions to the continuous-time reverse-diffusion trajectory.

  • Results

    3.78 FID on CIFAR-10 and 7.83 FID on ImageNet-64 are achieved with only one model evaluation.

  • Takeaways & Limitations

    Parallel decoding enables one-evaluation image generation and is presented as a step toward real-time diffusion-model sampling.

Abstract

from arXiv · show

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.

1. Introduction

Diffusion-model sampling is computationally slow because conventional approaches require many network evaluations or sequential trajectory steps. DSNO uses neural operators and Fourier-parameterized temporal convolutions to decode the reverse-diffusion trajectory in parallel, achieving strong one-evaluation results.

  • Diffusion-model sampling requires hundreds of neural-network evaluations, making it much slower than other generative models.
  • 12?

2. Background

The paper frames diffusion models in continuous time, where sampling solves a probability flow ODE backward from Gaussian noise. It also introduces Fourier neural operators as learned mappings between function spaces for efficiently solving differential equations.

  • Score-based generative models: Score-based models are treated interchangeably with diffusion models within a unified continuous-time framework.The forward diffusion process starts from the data distribution and ends in a standard Gaussian distribution.
  • Score-based generative models: The probability flow ODE produces the same marginal distributions as the corresponding diffusion process.This provides a deterministic differential-equation formulation of the sampling trajectory.
  • Score-based generative models: Sampling solves the probability flow ODE backward from time T to 0, starting from x(T).Numerical solvers include Euler’s method, multistep methods, and Heun’s second method.
  • Fourier neural operator: Fourier neural operators learn parametric maps between Banach spaces and can accelerate conventional PDE solvers.They are constructed from kernel integration layers whose kernels are parameterized by learnable weights.
  • Fourier neural operator: A Fourier neural operator lifts an input function, applies global Fourier-space integral layers with nonlinearities, and projects the result pointwise.The operator combines pointwise neural operators, residual connections, and trainable Fourier-space kernel parameters.

3. Learning the trajectory with neural operator

DSNO learns a neural operator that maps Gaussian initial conditions to the continuous probability-flow trajectory, enabling one-call trajectory prediction. Its Fourier temporal convolutions model temporal structure while preserving efficient parallel decoding.

  • Problem formulation: DSNO learns a mapping from Gaussian initial condition x(T) ∼ N(0, I) to the probability-flow trajectory over continuous time.The target is the solution operator of the diffusion ODE.
  • Problem formulation: Neural-operator theory establishes that the proposed architecture can approximate the diffusion ODE solution map and output the trajectory in one model call.
  • Temporal convolution: The temporal convolution block applies a Fourier convolution operator followed by a pointwise nonlinearity and an identity shortcut.The shortcut replaces a learned pointwise linear operator and is intended to preserve high-frequency information without extra cost.
  • Temporal convolution: The Fourier temporal convolution models trajectory correlations because its integral form structurally resembles the weighted integral form of the diffusion ODE solution.
  • Architecture and training: DSNO adds Fourier temporal convolution layers to each level of an existing diffusion-model U-Net, while the backbone processes temporal locations like batch elements.
  • Architecture and training: DSNO trains as an operator-learning problem by minimizing a weighted integral of trajectory error, with target trajectories generated by an existing solver or sampling method.
  • Parallel decoding: Fourier coefficients and inverse transforms allow outputs at all temporal locations to be generated in parallel, producing the continuous trajectory with one model evaluation.This parallel decoding relies on conditional independence of solutions at different times given the initial condition.
  • Parallel decoding: The method exploits the compact energy spectrum of diffusion ODE trajectories, where high-frequency temporal modes contribute little to the learning objective.

4. Experiments

Experiments evaluate DSNO’s sampling quality, speed, trajectory fidelity, and design choices across CIFAR-10 and ImageNet-64. DSNO achieves strong one-evaluation performance, while temporal convolution, loss weighting, discretization, and resolution affect results.

  • Sampling quality and speed: DSNO outperforms fast-sampling baselines with one model evaluation and achieves better FID than 2-step progressive distillation on CIFAR-10.The comparison uses FID and number of function evaluations.
  • Sampling quality and speed: DSNO achieves comparable FID to 2-step progressive distillation while outperforming the 1-step model on class-conditional ImageNet-64 with one model evaluation.Its speedup over progressive distillation is 1.7 times, and recall is comparable to ADM’s.
  • Trajectory prediction and reconstruction: DSNO preserves the original diffusion model’s trajectory and Gaussian-noise-to-image mapping for matched random seeds.Its predicted trajectory highly matches the original ODE solver trajectory on ImageNet-64.
  • Temporal convolution block: Temporal convolution is crucial to DSNO’s performance because its kernel-integration structure provides a better inductive bias for modeling trajectories in time.The comparison keeps temporal resolution, discretization, loss weighting, and batch size fixed.
  • Temporal resolution: Higher temporal resolution improves FID but increases computation cost; resolution 4 is favored because increasing from 4 to 8 gives only marginal benefit.The marginal benefit is attributed to the compact trajectory spectrum.

5. Related work

Related work divides fast diffusion sampling into training-free numerical solvers and training-based learned surrogates. Neural operators provide the function-space learning perspective motivating DSNO, while prior training-based methods include distillation and hybrid generative models.

  • ODE-based sampling: ODE-based samplers reduce diffusion sampling steps by exploiting structures such as semi-linearity and exponentially weighted integrals.Existing ODE-based methods reduce discretization steps to roughly 10–50.
  • Operator learning for solving PDEs: Neural operators learn mappings between function spaces and offer discretization invariance and universal approximation properties for scientific computing.These properties motivate DSNO’s Fourier temporal convolution design.
  • Training-based sampling: Training-based diffusion samplers replace parts or all of numerical solvers using approaches such as knowledge distillation, learned schedules, reverse covariance, or hybrid generative models.Progressive distillation reaches 4–8 steps with little quality loss, but knowledge-distillation methods face quality issues at one evaluation.

6. Conclusion and discussion

DSNO maps a Gaussian initial condition to the continuous-time reverse diffusion trajectory using temporal convolutions, achieving strong FID scores with one model evaluation.

  • DSNO maps a Gaussian initial condition to the continuous-time solution trajectory of the reverse diffusion process.
  • Temporal convolution layers model temporal correlations along the diffusion trajectory within the diffusion model backbone.
  • 3.78 FID on CIFAR-10 and 7.83 FID on ImageNet-64 are achieved with only one model evaluation.
  • The method is presented as a step toward real-time diffusion-model sampling for time-sensitive applications.

A.1. Energy spectrum

The energy spectrum characterizes trajectory frequency content through the Fourier transform and its conjugate; CIFAR-10 ODE trajectories concentrate most power below 5 Hz.

  • The discrete-time Fourier transform represents the frequency content of a periodic signal x(t) with period T.
  • The frequency mode indexes spectral components, with the time step defined as ∆=1/N.
  • The spectrum is defined as the product of x's Fourier transform and its complex conjugate.
  • Most power in DDPM++ cont. (VP) CIFAR-10 ODE trajectories lies below 5 Hz.

A.2. Background: neural operators

Neural operators learn nonlinear maps between function spaces from paired data, using kernel integration layers whose kernels are parameterized by learnable weights.

  • Neural operators learn a nonlinear map G from input functions in Banach space A to output functions in Banach space U.
  • The operator Gϕ is trained from paired samples by minimizing empirical risk.
  • Neural-operator architectures stack kernel integration layers with kernels parameterized by learnable weights.
  • Fourier neural operator is highlighted as an efficient neural-operator architecture for scientific computing problems involving PDEs.

A.3. Extended set of generated samples

The appendix provides an extended set of randomly generated ImageNet-64 samples from the DSNO model, illustrated in Figure 5.

  • The paper provides an extended set of randomly generated samples from the ImageNet-64 model.
  • Figure 5 displays random samples generated by DSNO on ImageNet-64.

A.4. Generalization to different resolution

DSNO generalizes from training temporal resolution 4 to predicting trajectories at resolution 8 on ImageNet-64, although the predicted trajectories are not perfectly smooth. The paper also identifies guided sampling and faster data collection as areas for further improvement.

  • Generalization to different resolution: DSNO predicts ImageNet-64 trajectories at temporal resolution 8 after training at resolution 4, demonstrating generalization to unseen time resolutions.The resulting trajectories are not perfectly smooth.
  • Future work: Guided diffusion sampling remains a more challenging setting that the paper leaves for future DSNO adaptation.The authors identify accelerating guided sampling as an interesting next step.
  • Future work: Advanced numerical solvers such as DPM solver can reduce the computation required to generate DSNO training data.The paper primarily uses DDIM for fair comparison but notes that DPM solver can speed up training-data generation.
Loading 2211.13449v3…