Umgekehrte Brückenübereinstimmungsdarstellung

Inverse Bridge Matching Distillation

February 3, 2025

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

cs.AI

Zusammenfassung

Das Erlernen von Diffusionsbrückenmodellen ist einfach; sie schnell und praktisch zu machen ist eine Kunst. Diffusionsbrückenmodelle (DBMs) sind eine vielversprechende Erweiterung von Diffusionsmodellen für Anwendungen in der Bild-zu-Bild-Übersetzung. Allerdings leiden DBMs wie viele moderne Diffusions- und Flussmodelle unter dem Problem langsamer Inferenz. Um dem entgegenzuwirken, schlagen wir eine neuartige Destillationstechnik auf Basis der inversen Brückenabgleichsformulierung vor und leiten das handhabbare Ziel ab, um es in der Praxis zu lösen. Im Gegensatz zu zuvor entwickelten DBM-Destillationstechniken kann die vorgeschlagene Methode sowohl bedingte als auch unbedingte Arten von DBMs destillieren, Modelle in einem Ein-Schritt-Generator destillieren und nur die korrupten Bilder für das Training verwenden. Wir evaluieren unseren Ansatz sowohl für bedingte als auch unbedingte Arten des Brückenabgleichs in einer breiten Palette von Setups, einschließlich Superresolution, JPEG-Restaurierung, Skizze-zu-Bild und anderen Aufgaben, und zeigen, dass unsere Destillationstechnik es uns ermöglicht, die Inferenz von DBMs um das 4- bis 100-fache zu beschleunigen und je nach speziellem Setup sogar eine bessere Generierungsqualität als das verwendete Lehrermodell zu bieten.

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.