Destilación de Emparejamiento Inverso de Puentes

Inverse Bridge Matching Distillation

February 3, 2025

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

cs.AI

Resumen

Aprender a modelar puentes de difusión es fácil; hacerlos rápidos y prácticos es un arte. Los modelos de puentes de difusión (DBMs, por sus siglas en inglés) son una extensión prometedora de los modelos de difusión para aplicaciones en traducción de imagen a imagen. Sin embargo, al igual que muchos modelos modernos de difusión y flujo, los DBMs sufren del problema de la inferencia lenta. Para abordarlo, proponemos una novedosa técnica de destilación basada en la formulación de coincidencia de puentes inversos y derivamos el objetivo viable para resolverlo en la práctica. A diferencia de las técnicas de destilación de DBM desarrolladas previamente, el método propuesto puede destilar tanto tipos condicionales como incondicionales de DBMs, destilar modelos en un generador de un solo paso y utilizar solo las imágenes corruptas para el entrenamiento. Evaluamos nuestro enfoque tanto para tipos condicionales como incondicionales de coincidencia de puentes en una amplia variedad de configuraciones, incluyendo superresolución, restauración de JPEG, boceto a imagen y otras tareas, y demostramos que nuestra técnica de destilación nos permite acelerar la inferencia de DBMs de 4x a 100x e incluso proporcionar una mejor calidad de generación que el modelo profesor utilizado, dependiendo de la configuración particular.

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.