Consistency Diffusion Bridge Models

TL;DR

This paper introduces a novel method for diffusion bridge models that significantly accelerates sampling speed and improves quality by learning a consistency function for the probability-flow ODE, enabling practical deployment.

Contribution

It proposes a new approach to speed up diffusion denoising bridge models by learning a consistency function for the PF-ODE, with flexible training paradigms and broad applicability.

Findings

Sampling speed improved by 4 to 50 times.

Produced higher quality images at same step compared to base DDBM.

Supported downstream tasks like semantic interpolation.

Abstract

Diffusion models (DMs) have become the dominant paradigm of generative modeling in a variety of domains by learning stochastic processes from noise to data. Recently, diffusion denoising bridge models (DDBMs), a new formulation of generative modeling that builds stochastic processes between fixed data endpoints based on a reference diffusion process, have achieved empirical success across tasks with coupled data distribution, such as image-to-image translation. However, DDBM's sampling process typically requires hundreds of network evaluations to achieve decent performance, which may impede their practical deployment due to high computational demands. In this work, inspired by the recent advance of consistency models in DMs, we tackle this problem by learning the consistency function of the probability-flow ordinary differential equation (PF-ODE) of DDBMs, which directly predicts the solution at a starting step given any point on the ODE trajectory. Based on a dedicated general-form ODE solver, we propose two paradigms: consistency bridge distillation and consistency bridge training, which is flexible to apply on DDBMs with broad design choices. Experimental results show that our proposed method could sample $4\times$ to $50\times$ faster than the base DDBM and produce better visual quality given the same step in various tasks with pixel resolution ranging from $64 \times 64$ to $256 \times 256$, as well as supporting downstream tasks such as semantic interpolation in the data space.