An Ordinary Differential Equation Sampler with Stochastic Start for Diffusion Bridge Models

TL;DR

This paper introduces a high-order ODE sampler with a stochastic start for diffusion bridge models, significantly improving inference speed and quality in image restoration and translation tasks without extra training.

Contribution

It proposes a novel stochastic start approach combined with Heun's solver for diffusion bridge models, reducing discretization errors and inference time.

Findings

Outperforms state-of-the-art in visual quality and FID scores

Reduces neural function evaluations significantly

Compatible with pretrained diffusion bridge models

Abstract

Diffusion bridge models have demonstrated promising performance in conditional image generation tasks, such as image restoration and translation, by initializing the generative process from corrupted images instead of pure Gaussian noise. However, existing diffusion bridge models often rely on Stochastic Differential Equation (SDE) samplers, which result in slower inference speed compared to diffusion models that employ high-order Ordinary Differential Equation (ODE) solvers for acceleration. To mitigate this gap, we propose a high-order ODE sampler with a stochastic start for diffusion bridge models. To overcome the singular behavior of the probability flow ODE (PF-ODE) at the beginning of the reverse process, a posterior sampling approach was introduced at the first reverse step. The sampling was designed to ensure a smooth transition from corrupted images to the generative trajectory while reducing discretization errors. Following this stochastic start, Heun's second-order solver is applied to solve the PF-ODE, achieving high perceptual quality with significantly reduced neural function evaluations (NFEs). Our method is fully compatible with pretrained diffusion bridge models and requires no additional training. Extensive experiments on image restoration and translation tasks, including super-resolution, JPEG restoration, Edges-to-Handbags, and DIODE-Outdoor, demonstrated that our sampler outperforms state-of-the-art methods in both visual quality and Frechet Inception Distance (FID).