Inverse Bridge Matching Distillation

Nikita Gushchin, David Li, Daniil Selikhanovych, Evgeny Burnaev, Dmitry Baranchuk, Alexander Korotin

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.

Lay Summary:

Modern AI systems that transform one image into another — for example, sharpening a blurry photo or turning a sketch into a realistic image — often rely on a class of tools called diffusion models. These models produce high-quality results but are painfully slow, sometimes taking hundreds or even thousands of steps to generate a single image.

Our research focuses on a special kind of these models called Diffusion Bridge Models (DBMs), which are well-suited for image-to-image tasks but face the same issue: slow generation speed. We introduce a new technique called Inverse Bridge Matching Distillation (IBMD) that significantly accelerates these models — making them up to 100 times faster — without sacrificing image quality.

Unlike earlier methods, our approach works universally for different types of DBMs and can even compress them into a single-step generator. We tested IBMD on a wide range of tasks, like super-resolution, inpainting, and image restoration, and found it not only faster but often better than the original models. This opens the door to making powerful AI-based image editing tools much more practical and accessible in everyday applications.

Primary Area: Deep Learning->Generative Models and Autoencoders

Keywords: Diffusion Bridges, Bridge Matching, Diffusion Bridges Distillation, Image Translation, Inverse Bridge Matching

Submission Number: 10865