Обратная дистилляция сопоставления мостов

Inverse Bridge Matching Distillation

February 3, 2025

Авторы: Nikita Gushchin, David Li, Daniil Selikhanovych, Evgeny Burnaev, Dmitry Baranchuk, Alexander Korotin

cs.AI

Аннотация

Обучение моделей диффузионных мостов несложно; сделать их быстрыми и практичными - это искусство. Модели диффузионных мостов (DBM) представляют собой многообещающее расширение моделей диффузии для применения в задачах перевода изображений. Однако, как и многие современные модели диффузии и потока, DBM страдают от проблемы медленного вывода. Для ее решения мы предлагаем новый метод дистилляции на основе формулировки обратного соответствия моста и выводим осуществимую цель для ее практического решения. В отличие от ранее разработанных методов дистилляции DBM, предложенный метод может дистиллировать как условные, так и безусловные типы DBM, дистиллировать модели в одношаговом генераторе и использовать только поврежденные изображения для обучения. Мы оцениваем наш подход как для условных, так и для безусловных типов соответствия мостов на широком наборе настроек, включая суперразрешение, восстановление JPEG, перевод скетчей в изображения и другие задачи, и показываем, что наш метод дистилляции позволяет ускорить вывод DBM с 4x до 100x и даже обеспечить лучшее качество генерации, чем используемая учителем модель, в зависимости от конкретной настройки.

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.