Bidirectional Diffusion Bridge Models

TL;DR

The paper introduces Bidirectional Diffusion Bridge Models (BDBM), a scalable single-network approach for efficient bidirectional image-to-image translation, outperforming existing unidirectional models and reducing computational costs.

Contribution

BDBM enables bidirectional translation using one network by leveraging the Chapman-Kolmogorov Equation, with analytical kernels for Gaussian marginals, connecting to existing bridge methods.

Findings

BDBM achieves bidirectional I2I translation with minimal extra cost.

It outperforms state-of-the-art bridge models in high-resolution tasks.

Analytical transition kernels are derived for Gaussian distributions.

Abstract

Diffusion bridges have shown potential in paired image-to-image (I2I) translation tasks. However, existing methods are limited by their unidirectional nature, requiring separate models for forward and reverse translations. This not only doubles the computational cost but also restricts their practicality. In this work, we introduce the Bidirectional Diffusion Bridge Model (BDBM), a scalable approach that facilitates bidirectional translation between two coupled distributions using a single network. BDBM leverages the Chapman-Kolmogorov Equation for bridges, enabling it to model data distribution shifts across timesteps in both forward and backward directions by exploiting the interchangeability of the initial and target timesteps within this framework. Notably, when the marginal distribution given endpoints is Gaussian, BDBM's transition kernels in both directions possess analytical forms, allowing for efficient learning with a single network. We demonstrate the connection between BDBM and existing bridge methods, such as Doob's h-transform and variational approaches, and highlight its advantages. Extensive experiments on high-resolution I2I translation tasks demonstrate that BDBM not only enables bidirectional translation with minimal additional cost but also outperforms state-of-the-art bridge models. Our source code is available at [https://github.com/kvmduc/BDBM||https://github.com/kvmduc/BDBM].