反向橋接匹配蒸餾

Inverse Bridge Matching Distillation

February 3, 2025

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

cs.AI

摘要

學習擴散橋模型相對容易；使其快速且實用則是一門藝術。擴散橋模型（DBMs）是擴散模型的一個有前途的延伸，適用於影像到影像的轉換應用。然而，像許多現代擴散和流模型一樣，DBMs 遭受緩慢推論的問題。為了解決這個問題，我們提出了一種基於逆橋匹配公式的新型蒸餾技術，並推導出可行的目標以實際解決它。與先前開發的 DBM 蒸餾技術不同，所提出的方法可以蒸餾有條件和無條件類型的 DBMs，蒸餾模型在一步生成器中，並僅使用損壞的影像進行訓練。我們在廣泛的設置中評估我們的方法，包括超分辨率、JPEG 恢復、素描到影像等任務，並展示我們的蒸餾技術使我們能夠將 DBMs 的推論加速從 4 倍到 100 倍，甚至根據特定設置提供比使用的教師模型更好的生成質量。

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.