Distilling Parallel Gradients for Fast ODE Solvers of Diffusion Models

TL;DR

This paper introduces EPD, a parallel gradient-based ODE solver for diffusion models that reduces sampling latency and maintains high image quality by leveraging multiple parallel gradient evaluations and a distillation training approach.

Contribution

The paper proposes a novel parallel gradient ODE solver, EPD, that improves sampling speed and quality in diffusion models through parallelization and learnable parameter distillation.

Findings

EPD achieves state-of-the-art FID scores at low latency levels.

EPD's parallel gradient approach reduces truncation errors effectively.

The method can be integrated into existing ODE samplers as a plugin.

Abstract

Diffusion models (DMs) have achieved state-of-the-art generative performance but suffer from high sampling latency due to their sequential denoising nature. Existing solver-based acceleration methods often face image quality degradation under a low-latency budget. In this paper, we propose the Ensemble Parallel Direction solver (dubbed as \ours), a novel ODE solver that mitigates truncation errors by incorporating multiple parallel gradient evaluations in each ODE step. Importantly, since the additional gradient computations are independent, they can be fully parallelized, preserving low-latency sampling. Our method optimizes a small set of learnable parameters in a distillation fashion, ensuring minimal training overhead. In addition, our method can serve as a plugin to improve existing ODE samplers. Extensive experiments on various image synthesis benchmarks demonstrate the effectiveness of our \ours~in achieving high-quality and low-latency sampling. For example, at the same latency level of 5 NFE, EPD achieves an FID of 4.47 on CIFAR-10, 7.97 on FFHQ, 8.17 on ImageNet, and 8.26 on LSUN Bedroom, surpassing existing learning-based solvers by a significant margin. Codes are available in https://github.com/BeierZhu/EPD.