PaperlayerView exact revision on arXiv

Source-linked AI summary

GENIE: Higher-Order Denoising Diffusion Solvers

Tim Dockhorn, Arash Vahdat, Karsten Kreis

arXiv:2210.05475v1stat.MLcs.LG

TL;DR

DDM synthesis requires solving a difficult generative differential equation efficiently. GENIE derives a higher-order Taylor solver using higher-order score information, distills the needed JVPs into a small prediction head, and achieves state-of-the-art solver performance while retaining encoding and guided sampling.

  • Problem

    DDM generative ODEs are difficult to solve because of complex score functions, creating a need for efficient samplers for fast synthesis.

  • Method

    GENIE applies a second-order truncated Taylor method and distills automatically differentiated higher-order JVP terms into a small neural-network head on the first-order score model.

  • Results

    GENIE achieves state-of-the-art performance on multiple image-modeling benchmarks and outperforms previous solvers and samplers for DDM generative differential equations.

  • Takeaways & Limitations

    Because GENIE solves the true generative ODE, it retains applications such as image encoding, interpolation, and guided sampling.

  • Takeaways & Limitations

    GENIE remains slightly slower than approaches that abandon the DDM differential-equation framework altogether.

Abstract

from arXiv · show

Denoising diffusion models (DDMs) have emerged as a powerful class of generative models. A forward diffusion process slowly perturbs the data, while a deep model learns to gradually denoise. Synthesis amounts to solving a differential equation (DE) defined by the learnt model. Solving the DE requires slow iterative solvers for high-quality generation. In this work, we propose Higher-Order Denoising Diffusion Solvers (GENIE): Based on truncated Taylor methods, we derive a novel higher-order solver that significantly accelerates synthesis. Our solver relies on higher-order gradients of the perturbed data distribution, that is, higher-order score functions. In practice, only Jacobian-vector products (JVPs) are required and we propose to extract them from the first-order score network via automatic differentiation. We then distill the JVPs into a separate neural network that allows us to efficiently compute the necessary higher-order terms for our novel sampler during synthesis. We only need to train a small additional head on top of the first-order score network. We validate GENIE on multiple image generation benchmarks and demonstrate that GENIE outperforms all previous solvers. Unlike recent methods that fundamentally alter the generation process in DDMs, our GENIE solves the true generative DE and still enables applications such as encoding and guided sampling. Project page and code: https://nv-tlabs.github.io/GENIE.

1 Introduction

DDM synthesis requires solving a difficult generative differential equation efficiently. GENIE introduces a higher-order solver that uses higher-order score information and preserves the true generative process while accelerating sampling.

  • DDM generative ODEs are difficult to solve because their score functions are complex, motivating efficient tailored samplers.
  • GENIE achieves state-of-the-art performance on multiple image-modeling benchmarks while solving the true generative ODE.
  • Because GENIE retains the generative differential-equation framework, it supports image encoding and guided sampling.
  • GENIE is a second-order ODE solver that uses truncated Taylor methods to improve DDM sampling efficiency.
  • GENIE extracts required higher-order terms from the first-order score model using automatic differentiation and avoids finite-difference approximations.
  • The method distills Jacobian-vector products into a small neural network head for efficient synthesis-time computation.

2 Background

Continuous-time DDMs define forward and reverse stochastic processes, with generation also expressible through a Probability Flow ODE. Sampling uses a learned score model to numerically solve this generative dynamics, with DDIM providing a fast Euler-based solver.

  • Continuous-time DDMs diffuse data samples over time using functions that control signal and noise levels.
  • The forward diffusion can be represented as a variance-preserving stochastic differential equation driven by a standard Wiener process.
  • A reverse-time generative SDE approximately inverts the forward diffusion process.
  • The Probability Flow ODE is marginally equivalent to the reverse generative SDE and uses the score function of the perturbed data distribution.
  • Generation starts from Gaussian noise and numerically simulates the Probability Flow ODE or generative SDE using a learned score model.
  • DDIM is a popular fast solver because it is Euler’s method applied to an ODE obtained through re-parameterization.

3 Higher-Order Denoising Diffusion Solver

GENIE applies truncated Taylor methods to the DDIM ODE, using higher-order derivatives to capture local curvature and permit larger solver steps. Because direct automatic differentiation is costly, the required Jacobian-vector products are distilled into a small prediction head for efficient synthesis.

  • Higher-Order Solver: GENIE applies the second truncated Taylor method to the DDIM ODE, extending Euler’s method with higher-order gradient terms that model local curvature.The resulting extrapolation is quadratic in time and is intended to support larger time steps.
  • Toy Distribution: 25 solver steps with the analytical score function produce DDIM samples that enter low-density regions, while GENIE samples more closely resemble the ground-truth distribution.The toy experiment uses a known analytical score function and compares both methods against the ground-truth distribution.
  • Derivative Estimation: Finite-difference derivatives approximate the analytical derivative poorly at small t, where the score function becomes complex.The comparison measures the mean L2-norm between the analytical derivative and its first-order forward-difference approximation.
  • Derivative Estimation: The derivative needed by GENIE decomposes into two Jacobian-vector products and one simpler time-derivative term.This decomposition avoids explicitly forming the full high-dimensional Jacobian of the score network.
  • Gradient Distillation: Direct automatic differentiation would make each GENIE step at least twice as costly as DDIM because its forward passes cannot be parallelized.The vector part of the first JVP depends on the score-network output, creating sequential overhead during sampling.
  • Gradient Distillation: GENIE distills the required derivatives into a small prediction head that reuses the first-order score model’s features, with less than 2% overhead on CIFAR-10.The head receives internal features, the time embedding, the current sample, and the score output; training uses automatic-differentiation targets.
  • Gradient Distillation: GENIE directly uses higher-order scores through distilled JVPs rather than finite-difference or other approximations.The authors identify this as the first DDM generative solver to use higher-order scores explicitly in this form.

4 Related Work

Prior work accelerates diffusion sampling by changing timesteps, solvers, diffusion formulations, latent representations, or learned samplers. GENIE belongs to the higher-order ODE-solver line while retaining the original generative differential-equation framework.

  • Accelerated Sampling from DDMs: Prior diffusion-sampling methods reduce or readjust timesteps, optimize reverse variances, or use ODE and SDE solvers such as DDIM and Runge–Kutta methods.These approaches target faster synthesis through discretization, variance design, or numerical integration.
  • Accelerated Sampling from DDMs: Learning-based approaches accelerate sampling by optimizing diffusion parameters, distilling DDIM into student models, or replacing Gaussian samplers with generative adversarial networks.Some distilled approaches enable sampling in as few as a single step.
  • Accelerated Sampling from DDMs: Other methods change the diffusion process or train diffusion models in variational-autoencoder latent spaces, and GENIE is complementary to them.
  • Higher-Order ODE Gradients beyond DDMs: Higher-order truncated Taylor methods and related higher-order-gradient techniques are also used for neural ODE solvers, regularization, and stiff or non-stiff systems.
  • Concurrent Works: Concurrent work applies exponential-integrator, Runge–Kutta, linear-multistep, semi-linear, and specialized solvers to diffusion-model differential equations.These works also introduce step-size schedulers or deterministic and stochastic solver variants.

5 Experiments

GENIE is evaluated across image-generation benchmarks, guidance, encoding, and upsampling tasks. It generally improves few-step synthesis while preserving diffusion-ODE capabilities such as guidance and image encoding.

  • Evaluation setup: Experiments cover CIFAR-10, LSUN Bedrooms, LSUN Church-Outdoor, ImageNet, and Cats, with quality measured using FID.The Cats experiments use a 128 × 128 →512 × 512 diffusion upsampler.
  • Implementation and overhead: The prediction head adds 1.47%, 2.83%, 14.0%, and 14.4% computational overhead on CIFAR-10, ImageNet, LSUN Bedrooms, and LSUN Church-Outdoor, respectively.The experiments account for this overhead by adjusting the reported NFEs.
  • Image generation: GENIE outperforms same-checkpoint baselines for NFEs ∈{10, 15, 20, 25} across all four benchmark datasets.At NFEs=5, GENIE also leads on CIFAR-10 and ImageNet, while DDIM is slightly better on the LSUN datasets.
  • Image generation: 13.1 FID is achieved by adaptive Runge–Kutta 4(5) at 38 NFEs, illustrating the low-NFE setting targeted by GENIE.This comparison uses the same checkpoint and the ProbabilityFlow ODE.
  • Image generation: GENIE’s higher-order gradient information can be efficiently distilled: automatic-differentiation derivatives achieve FIDs of 9.22, 4.11, 3.54, and 3.46 at 10, 20, 30, and 40 NFEs.The reported NFE values account for an assumed 100% additional computational overhead.
  • Guidance and encoding: GENIE clearly outperforms DDIM for classifier-free guidance on ImageNet, especially at few NFEs, while also producing high-quality guided images.The guidance formulation combines conditional and unconditional score models through a guidance scale.
  • Guidance and encoding: GENIE reconstructs images more accurately than DDIM in encode-decode experiments, using the generative ODE in reverse for encoding.The comparison evaluates reconstructions across different NFEs used for both encoding and decoding.
  • Upsampling: GENIE outperforms all other evaluated upsampler methods at NFEs=15 while using only five NFEs.The upsampler is evaluated in isolation on Cats and produces 128 × 128 →512 × 512 outputs.

6 Conclusions

GENIE applies truncated Taylor methods to construct a higher-order solver for DDM generative ODEs and distills the needed derivative into a small prediction head. The approach improves upon DDIM by capturing local curvature, while remaining within the differential-equation framework and retaining guided sampling and encoding capabilities.

  • Method: GENIE is a higher-order ODE solver for DDMs based on truncated Taylor methods.The second-order method is applied to the approximate Probability Flow or DDIM ODE.
  • Method: GENIE improves upon DDIM by capturing local curvature in the ODE’s gradient field, enabling larger integration steps.The prediction head provides the higher-order derivatives needed for this solver.
  • Method: The method uses the approximate generative DDIM ODE obtained by replacing the unknown score with a learned score model.The underlying probability-flow formulation uses the score function of the perturbed data distribution.
  • Method: The solver distills the derivative dγtϵθ into a small neural network kψ, using automatic differentiation during training and efficient network queries during inference.This avoids repeatedly computing the derivative through automatic differentiation at synthesis time.

B.1 Theoretical Bounds for the Truncated Taylor Method

The truncated Taylor method (TTM) reduces local truncation error by incorporating higher-order derivatives, motivating higher-order solvers for diffusion-model ODEs. GENIE applies this principle to the DDIM ODE, while its derivative approximation behaves differently for the Probability Flow ODE.

  • Theoretical bounds: The p-th TTM approximates an ODE solution with a Taylor expansion whose local truncation error is proportional to h_n^(p+1) for small h_n.Higher orders therefore reduce error when the solver step is small.
  • Theoretical bounds: Higher-order TTM methods are preferable for accurately solving diffusion-model ODEs such as the Probability Flow ODE and DDIM ODE.
  • DDIM ODE: Under the ideal derivative trick, the second TTM applied to the DDIM ODE reduces to Euler’s method because the relevant derivative term cancels to zero.This provides an approximate second-order interpretation of the DDIM solver.
  • Probability Flow ODE: The same ideal derivative trick does not produce the corresponding cancellation for the Probability Flow ODE, so Euler’s method is not an approximate second-order solver there.
  • Higher-order experiments: The third TTM improves single-step error over the second TTM on the 2D toy distribution when analytical score derivatives are available.For CIFAR-10, however, the second TTM performs better than the third TTM with 15 or fewer solver steps, potentially because the learned score model and its higher-order derivatives are inaccurate or unregularized.
  • GENIE construction: GENIE combines a second TTM with a learned prediction head because exact automatic-differentiation derivatives would require additional backward passes during inference.If the prediction head matches the ground-truth derivatives, GENIE recovers the exact second TTM.

C.2.1 Model Architecture

GENIE uses a small prediction head built on the first-order score network to estimate the higher-order terms needed by its sampler while limiting computational overhead.

  • Prediction-head architecture: The prediction head uses a modified BigGAN residual-block architecture with a single residual block to minimize computational overhead.
  • Prediction-head architecture: The head receives the last score-network feature layer, x_t, the predicted noise ϵ_θ(x_t,t), and a time embedding as inputs.For the upsampler, it also conditions on the noisy up-scaled lower-resolution image.
  • Prediction-head architecture: The paper reports the prediction-head hyperparameters and their additional computational overhead in Table 5.

C.2.2 Training Details

GENIE trains prediction heads to reproduce higher-order derivatives computed from the score network, then uses those heads during sampling instead of costly automatic differentiation.

  • Training details: Mixed-precision training produced NaNs when automatically differentiating the time derivative at small t, so the authors disabled mixed precision.They note that disabling it only for this derivative could improve training performance.
  • Parameterization: The prediction heads use a mixed parameterization intended to produce well-behaved outputs under an approximately Normal data assumption.The parameterization replaces analytic terms with neural-network outputs while relating the predicted noise to the normalized data scale.
  • Training and sampling: GENIE sampling supports arbitrary time discretizations, including linear and quadratic striding, and initializes samples from a standard Normal distribution.
  • Training and sampling: During training, GENIE computes the derivative d_γtϵ_θ with automatic differentiation and trains a prediction head k_ψ to approximate it.At inference, the learned head replaces the derivative calculation.
  • Computational cost: The prediction-head forward pass adds computational overhead because it follows the score-model forward pass, which the paper measures using repeated synchronized GPU timings.
  • Training and sampling: The derivative target contains spatial and temporal Jacobian-vector products, and the paper also considers a score-matching-like objective for the spatial JVP without automatic differentiation.The authors report that their automatic-differentiation distillation approach worked significantly better than this alternative.

E Toy Experiments

Toy experiments use a ground-truth distribution with analytically available score derivatives to compare DDIM and GENIE across solver-step budgets. GENIE and DDIM produce structured samples across the tested settings, including particularly at ten steps.

  • Setup: The toy experiments use a ground-truth distribution whose score functions and all derivatives can be computed analytically because it is a Gaussian mixture.
  • Sampling comparison: The authors compare DDIM and GENIE using the analytical score function with 25 solver steps, then repeat the comparison with 5, 10, 15, and 20 steps.
  • Sampling comparison: At n = 10, both DDIM and GENIE generate samples in interesting patterns on the complex 2D toy distribution.Figure 12 visualizes samples produced by both methods across solver-step counts.

F.1 Evaluation Metrics, Baselines, and Datasets

The evaluation uses FID and recall across several image datasets, with baselines and synthesis variants covering sampling, guidance, encoding, decoding, and interpolation.

  • Evaluation Metrics: Sample quality is measured with Fréchet Inception Distance, while sample diversity is evaluated with recall.Reference statistics use 50k training samples except for ImageNet, which uses its full training set, and Cats, which combines training and validation data.
  • Synthesis Strategy: The synthesis strategy includes an analytical first step that replaces the learned score with the analytical score of N(0, I) at the initial time.The resulting saved function evaluation is used as an additional ODE step, especially in the low-NFE regime.
  • Guided Sampling: Classifier-free guidance replaces the conditional score with a guidance-scaled combination of conditional and unconditional predictions.The same conditional prediction head is reused for the guided setting.
  • Encoding and Interpolation: Encoding runs the generative ODE backward from t = 10^-3 to t = 1, while decoding runs it forward and optionally denoises the result.The encode-decode setup supports spherical interpolation between encoded latent points.

F.6 Extended Quantitative Results

Extended experiments show that synthesis choices matter across NFE budgets, while GENIE achieves lower discretization error and strong diversity relative to fast-sampling baselines.

  • Synthesis Strategy: AFS generally improves results, particularly for NFEs ≤15, whereas denoising becomes more important at larger NFEs.For the Cats models, denoising initially hurt performance and was therefore not tested in all settings.
  • Recall Scores: GENIE achieves recall on par with F-PNDM for NFE≥15, while recall decreases for all methods as NFEs decrease.F-PNDM cannot run at NFE≤10 because of its additional Runge–Kutta warm-up iterations.
  • Striding Schedules: 13.9 to 11.2 FID is the improvement reported for GENIE at NFE=5 under the striding-schedule search.The quadratic schedule uses ρ = 2, while ρ = 1.5 and ρ = 2.5 are also tested for GENIE and DDIM.
  • Discretization Errors: GENIE outperforms DDIM, S-PNDM, and F-PNDM in global discretization error across all tested NFEs.Errors are measured against trajectories generated by DDIM with 1k NFEs using the same latent vectors and score model.
  • Discretization Errors: GENIE has smaller local truncation error than DDIM at starting times t ∈ {0.1, 0.2, 0.5}.The comparison uses one-step predictions and L2-distance in data space averaged over 100 predictions.

F.7 Extended Qualitative Results

Additional figures show GENIE samples with finer details, sharper contrast, and less blur than DDIM across several datasets, alongside Cats upsampling examples and broader method context.

  • Qualitative Comparisons: GENIE samples generally exhibit finer details, sharper contrast, and less blur than standard DDIM on Church-Outdoor, ImageNet, and Cats.The comparisons include Cats upsampler outputs conditioned on test-set images.
  • Cats Upsampling: The GENIE Cats end-to-end example uses 25 base-model function evaluations and five upsampler evaluations.An upsampler evaluation is roughly four times as expensive as a base-model evaluation.
  • Related Methods: GENIE uses a small prediction head on top of the first-order score predictor, unlike Bao et al., who learn sampling covariance matrices.The paper also distinguishes GENIE’s higher-order ODE gradients and deterministic sampling from Bao et al.’s stochastic approach.
  • Cats Upsampling: The GENIE upsampler generates 128 × 128 test-set image outputs using only five function evaluations.Additional figures show high-resolution Cats images generated from base-model and test-set samples.
  • Future Directions: The authors speculate that GENIE could be combined with Progressive Distillation to produce larger or more accurate distillation steps.They note that multistep methods are unsuitable for this scheme because previous predictions are unavailable during progressive-distillation training.
Loading 2210.05475v1…