Distillazione tramite corrispondenza inversa dei ponti.

Inverse Bridge Matching Distillation

February 3, 2025

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

cs.AI

Abstract

Apprendere i modelli di ponte di diffusione è facile; renderli veloci e pratici è un'arte. I modelli di ponte di diffusione (DBM) sono un'estensione promettente dei modelli di diffusione per applicazioni nella traduzione immagine-immagine. Tuttavia, come molti modelli moderni di diffusione e flusso, i DBM soffrono del problema dell'inferezza lenta. Per affrontarlo, proponiamo una nuova tecnica di distillazione basata sulla formulazione dell'abbinamento inverso del ponte e deriviamo l'obiettivo praticabile per risolverlo in pratica. A differenza delle tecniche di distillazione DBM sviluppate in precedenza, il metodo proposto può distillare sia tipi condizionali che incondizionali di DBM, distillare modelli in un generatore a un passo e utilizzare solo le immagini corrotte per l'addestramento. Valutiamo il nostro approccio per entrambi i tipi di abbinamento di ponti, condizionali e incondizionali, su un ampio set di configurazioni, inclusi super-risoluzione, ripristino JPEG, schizzo-immagine e altre attività, e mostriamo che la nostra tecnica di distillazione ci consente di accelerare l'inferezza dei DBM da 4x a 100x e persino fornire una migliore qualità di generazione rispetto al modello di riferimento utilizzato a seconda della configurazione specifica.

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.