Inverse Bridge Matching Distillation

TL;DR

This paper introduces a novel distillation method for diffusion bridge models that significantly accelerates inference speed from 4x to 100x while maintaining or improving quality across various image translation tasks.

Contribution

The proposed inverse bridge matching distillation technique is the first to distill both conditional and unconditional DBMs into a one-step generator using only corrupted images for training.

Findings

Inference speed improved 4x to 100x

Achieved better generation quality than teacher models in some setups

Effective across multiple image translation tasks

Abstract

Learning diffusion bridge models is easy; making them fast and practical is an art. Diffusion bridge models (DBMs) are a promising extension of diffusion models for applications in image-to-image translation. However, like many modern diffusion and flow models, DBMs suffer from the problem of slow inference. To address it, we propose a novel distillation technique based on the inverse bridge matching formulation and derive the tractable objective to solve it in practice. Unlike previously developed DBM distillation techniques, the proposed method can distill both conditional and unconditional types of DBMs, distill models in a one-step generator, and use only the corrupted images for training. We evaluate our approach for both conditional and unconditional types of bridge matching on a wide set of setups, including super-resolution, JPEG restoration, sketch-to-image, and other tasks, and show that our distillation technique allows us to accelerate the inference of DBMs from 4x to 100x and even provide better generation quality than used teacher model depending on particular setup. We provide the code at https://github.com/ngushchin/IBMD