역 다리 매칭 증류

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을 증류할 수 있으며, 한 단계 생성기에서 모델을 증류하고 오직 손상된 이미지만을 사용하여 훈련합니다. 우리는 초해상도, 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.