逆ブリッジマッチング蒸留

Inverse Bridge Matching Distillation

February 3, 2025

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

cs.AI

要旨

拡散ブリッジモデルの学習は簡単ですが、それを高速かつ実用的にすることは芸術です。拡散ブリッジモデル（DBM）は、画像間変換のアプリケーションにおける拡散モデルの有望な拡張です。しかし、多くの現代の拡散モデルやフローモデルと同様に、DBMには推論が遅いという問題があります。この問題に対処するために、逆ブリッジマッチングの形式に基づいた新しい蒸留技術を提案し、実践的に解決するための取り扱い可能な目的を導出します。従来開発されたDBM蒸留技術とは異なり、提案された方法は条件付きおよび無条件のDBMの両方を蒸留し、1ステップのジェネレータでモデルを蒸留し、訓練には破損した画像のみを使用します。私たちのアプローチを、超解像、JPEG復元、スケッチから画像へなどの幅広いセットアップで条件付きおよび無条件のブリッジマッチングについて評価し、蒸留技術によってDBMの推論を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.