Distillation de Correspondance Inverse de Pont

Inverse Bridge Matching Distillation

February 3, 2025

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

cs.AI

Résumé

Apprendre à modéliser les ponts de diffusion est facile ; les rendre rapides et pratiques relève de l'art. Les modèles de pont de diffusion (DBM) sont une extension prometteuse des modèles de diffusion pour des applications de traduction d'images à images. Cependant, comme de nombreux modèles de diffusion et de flux modernes, les DBM souffrent du problème d'inférence lente. Pour y remédier, nous proposons une nouvelle technique de distillation basée sur la formulation de correspondance de pont inverse et dérivons l'objectif traitable pour le résoudre en pratique. Contrairement aux techniques de distillation DBM développées précédemment, la méthode proposée peut distiller à la fois des types de DBM conditionnels et inconditionnels, distiller des modèles dans un générateur en une étape, et n'utiliser que les images corrompues pour l'entraînement. Nous évaluons notre approche pour les types de correspondance de pont conditionnels et inconditionnels sur un large ensemble de configurations, y compris la super-résolution, la restauration JPEG, le croquis à l'image, et d'autres tâches, et montrons que notre technique de distillation nous permet d'accélérer l'inférence des DBM de 4x à 100x et même de fournir une meilleure qualité de génération que le modèle enseignant utilisé en fonction de la configuration particulière.

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.