逆桥匹配蒸馏

Inverse Bridge Matching Distillation

February 3, 2025

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

cs.AI

摘要

学习扩散桥模型很容易；使其快速和实用则是一门艺术。扩散桥模型（DBMs）是扩散模型的一个有前途的延伸，可用于图像到图像的翻译应用。然而，像许多现代扩散和流模型一样，DBMs存在推断速度慢的问题。为了解决这个问题，我们提出了一种基于逆桥匹配公式的新型提炼技术，并推导出可行的目标以在实践中解决它。与先前开发的DBM提炼技术不同，所提出的方法可以提炼条件和无条件类型的DBMs，提炼模型在一个步骤生成器中，并且仅使用损坏的图像进行训练。我们在一系列设置中评估了我们的方法，包括超分辨率、JPEG 恢复、素描到图像等任务，并展示了我们的提炼技术使我们能够将DBMs的推断加速从4倍到100倍，甚至根据特定设置提供比使用的教师模型更好的生成质量。

English

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.