Destilação de Correspondência de Ponte Inversa

Inverse Bridge Matching Distillation

February 3, 2025

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

cs.AI

Resumo

Aprender a modelar pontes de difusão é fácil; torná-las rápidas e práticas é uma arte. Os modelos de ponte de difusão (DBMs) são uma extensão promissora dos modelos de difusão para aplicações em tradução de imagem para imagem. No entanto, assim como muitos modelos modernos de difusão e fluxo, os DBMs sofrem do problema de inferência lenta. Para abordá-lo, propomos uma técnica de destilação inovadora com base na formulação de correspondência de ponte inversa e derivamos o objetivo viável para resolvê-lo na prática. Ao contrário das técnicas de destilação de DBM desenvolvidas anteriormente, o método proposto pode destilar tanto tipos condicionais quanto incondicionais de DBMs, destilar modelos em um gerador de um passo e usar apenas imagens corrompidas para treinamento. Avaliamos nossa abordagem para ambos os tipos condicionais e incondicionais de correspondência de ponte em uma ampla variedade de configurações, incluindo super-resolução, restauração de JPEG, esboço para imagem e outras tarefas, e mostramos que nossa técnica de destilação nos permite acelerar a inferência de DBMs de 4x a 100x e até mesmo fornecer melhor qualidade de geração do que o modelo professor utilizado, dependendo da configuração específica.

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.