Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Universal Inverse Distillation for Matching Models with Real-Data Supervision (No GANs)

Nikita Kornilov, David Li, Tikhon Mavrin, Aleksei Leonov, Nikita Gushchin, Evgeny Burnaev, Iaroslav Koshelev, Alexander Korotin

TL;DR

The RealUID approach offers a simple theoretical foundation that covers previous distillation methods for Flow Matching and Diffusion models, and is also extended to their modifications, such as Bridge Matching and Stochastic Interpolants.

Abstract

While achieving exceptional generative quality, modern diffusion, flow, and other matching models suffer from slow inference, as they require many steps of iterative generation. Recent distillation methods address this by training efficient one-step generators under the guidance of a pre-trained teacher model. However, these methods are often constrained to only one specific framework, e.g., only to diffusion or only to flow models. Furthermore, these methods are naturally data-free, and to benefit from the usage of real data, it is required to use an additional complex adversarial training with an extra discriminator model. In this paper, we present RealUID, a universal distillation framework for all matching models that seamlessly incorporates real data into the distillation procedure without GANs. Our RealUID approach offers a simple theoretical foundation that covers previous distillation methods for Flow Matching and Diffusion models, and is also extended to their modifications, such as Bridge Matching and Stochastic Interpolants. The code can be found in https://github.com/David-cripto/RealUID.

Universal Inverse Distillation for Matching Models with Real-Data Supervision (No GANs)

TL;DR

The RealUID approach offers a simple theoretical foundation that covers previous distillation methods for Flow Matching and Diffusion models, and is also extended to their modifications, such as Bridge Matching and Stochastic Interpolants.

Abstract

While achieving exceptional generative quality, modern diffusion, flow, and other matching models suffer from slow inference, as they require many steps of iterative generation. Recent distillation methods address this by training efficient one-step generators under the guidance of a pre-trained teacher model. However, these methods are often constrained to only one specific framework, e.g., only to diffusion or only to flow models. Furthermore, these methods are naturally data-free, and to benefit from the usage of real data, it is required to use an additional complex adversarial training with an extra discriminator model. In this paper, we present RealUID, a universal distillation framework for all matching models that seamlessly incorporates real data into the distillation procedure without GANs. Our RealUID approach offers a simple theoretical foundation that covers previous distillation methods for Flow Matching and Diffusion models, and is also extended to their modifications, such as Bridge Matching and Stochastic Interpolants. The code can be found in https://github.com/David-cripto/RealUID.

Paper Structure

This paper contains 87 sections, 7 theorems, 95 equations, 15 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Backgrounds on training and distilling matching models
  3. Diffusion and flow models
  4. Universal loss for matching models
  5. Distillation of matching-based models
  6. GANs for real data incorporation
  7. Universal distillation of matching models with real data
  8. Universal Inverse Distillation
  9. Making loss tractable via linearization.
  10. Summary.
  11. Relation to prior distillation works
  12. Connection with Inverse Optimization
  13. RealUID: natural approach for real data incorporation
  14. Role of coefficients $\alpha, \beta$.
  15. Choice of coefficients $\alpha, \beta$.
  16. ...and 72 more sections

Key Result

Theorem 1

Let teacher $f^* := \mathop{\mathrm{arg\,min}}\limits_f \mathcal{L}_{\text{UM}}(f, p_0^*)$ be the minimizer of UM loss (Def. def: UM loss) on real data $p_0^* \in \mathcal{P}(\mathbb{R}^D)$. Then, real data generator $G_{\theta^*}$ s.t. $p_0^{\theta^*} = p_0^*$ is a solution to the min-max optimizat

Figures (15)

  • Figure 1: Pipeline of our RealUID distillation framework (§ \ref{['sec: unified view']}) with the direct incorporation of real data $p^*_0$ adjusted by parameters $\alpha, \beta \in (0,1]$. The figure depicts flow matching models predicting denoised samples. It distills a costly frozen teacher model $f^*$ (blue) into a one-step generator $G_\theta$ (red) upon min-max optimization of $\mathcal{L}^{\alpha, \beta}_{\text{R-UID}}(f, p_0^\theta)$ loss over fake model $f$ (green) and generator distribution $p_0^\theta$ with parameters $\theta$. It updates the fake model several times per one generator update for stability. Algorithm's pseudocode is located in Appendix \ref{['sec: M-UID for FM']}.
  • Figure 2: Evolution of FID during CIFAR-10 distillation for (i) the UID (FGM) baseline, (ii) the best-performing RealUID configurations, and (iii) subsequent fine-tuning, evaluated in both unconditional and conditional settings. The performances of Teacher Flow and UID+GAN are indicated by horizontal lines in their respective colors.
  • Figure 3: RealUID loss for $1D$-Gaussians under various coefficients $(\alpha, \beta)$.
  • Figure 4: Learned generators for RealUID loss between 1D-Gaussians with corrupted teachers.
  • Figure 5: Evolution of FID during unconditional CIFAR-10 distillation for the data-free SiD loss ($\alpha = \beta = 1.0$) and our RealSiD loss for $\alpha_{\text{SiD}} = 0.5$ (left) and $\alpha_{\text{SiD}} = 1.2$ (right).
  • ...and 10 more figures

Theorems & Definitions (16)

  • Definition 1
  • Theorem 1: Real data generator minimizes UID loss
  • Lemma 1: UID loss minimizes squared $\ell_2$-distance
  • Definition 2
  • Theorem 2: Real data generator minimizes RealUID loss
  • Lemma 2: Distance minimized by RealUID loss
  • Lemma 3: RealUID split form
  • proof
  • proof : Proof of Lemma \ref{['lem: M-UID distance']}.
  • Definition 3
  • ...and 6 more