Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distilling Two-Timed Flow Models by Separately Matching Initial and Terminal Velocities

Pramook Khungurn, Pratch Piyawongwisal, Sira Sriswasdi, Supasorn Suwajanakorn

TL;DR

This work introduces ITVM, a three-term loss for distilling a flow-matching model into a two-time flow model (TTFM) to enable fast, flexible sampling. ITVM combines initial velocity matching, an initial-average velocity term, and a terminal velocity consistency term, while using an EMA-smoothed model as the velocity target to stabilize training. By enforcing a Lagrangian PDE-inspired construction through velocity alignments and self-consistency, ITVM yields improved few-step generation across 2D, tabular, and image datasets compared to LFMD, PID, and EFMD. The approach provides practical guidance on hyperparameters (e.g., $\tau$, intermediate time $u$, EMA decay) and contributes to the broader goal of efficient diffusion-based generative modeling without sacrificing accuracy.

Abstract

A flow matching model learns a time-dependent vector field $v_t(x)$ that generates a probability path $\{ p_t \}_{0 \leq t \leq 1}$ that interpolates between a well-known noise distribution ($p_0$) and the data distribution ($p_1$). It can be distilled into a two-timed flow model (TTFM) $φ_{s,x}(t)$ that can transform a sample belonging to the distribution at an initial time $s$ to another belonging to the distribution at a terminal time $t$ in one function evaluation. We present a new loss function for TTFM distillation called the \emph{initial/terminal velocity matching} (ITVM) loss that extends the Lagrangian Flow Map Distillation (LFMD) loss proposed by Boffi et al. by adding redundant terms to match the initial velocities at time $s$, removing the derivative from the terminal velocity term at time $t$, and using a version of the model under training, stabilized by exponential moving averaging (EMA), to compute the target terminal average velocity. Preliminary experiments show that our loss leads to better few-step generation performance on multiple types of datasets and model architectures over baselines.

Distilling Two-Timed Flow Models by Separately Matching Initial and Terminal Velocities

TL;DR

This work introduces ITVM, a three-term loss for distilling a flow-matching model into a two-time flow model (TTFM) to enable fast, flexible sampling. ITVM combines initial velocity matching, an initial-average velocity term, and a terminal velocity consistency term, while using an EMA-smoothed model as the velocity target to stabilize training. By enforcing a Lagrangian PDE-inspired construction through velocity alignments and self-consistency, ITVM yields improved few-step generation across 2D, tabular, and image datasets compared to LFMD, PID, and EFMD. The approach provides practical guidance on hyperparameters (e.g., , intermediate time , EMA decay) and contributes to the broader goal of efficient diffusion-based generative modeling without sacrificing accuracy.

Abstract

A flow matching model learns a time-dependent vector field that generates a probability path that interpolates between a well-known noise distribution () and the data distribution (). It can be distilled into a two-timed flow model (TTFM) that can transform a sample belonging to the distribution at an initial time to another belonging to the distribution at a terminal time in one function evaluation. We present a new loss function for TTFM distillation called the \emph{initial/terminal velocity matching} (ITVM) loss that extends the Lagrangian Flow Map Distillation (LFMD) loss proposed by Boffi et al. by adding redundant terms to match the initial velocities at time , removing the derivative from the terminal velocity term at time , and using a version of the model under training, stabilized by exponential moving averaging (EMA), to compute the target terminal average velocity. Preliminary experiments show that our loss leads to better few-step generation performance on multiple types of datasets and model architectures over baselines.
Paper Structure (40 sections, 7 theorems, 55 equations, 9 figures, 6 tables)

This paper contains 40 sections, 7 theorems, 55 equations, 9 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Previous Works
  3. Preliminary
  4. Method
  5. Experiments
  6. Comparison with Other Losses
  7. Analysis of Design Choices and Hyperparameters
  8. Conclusion
  9. Proofs
  10. Metrics Used for Performance Comparison
  11. KL Divergence
  12. Preliminary
  13. Problem setup
  14. KL Divergence as Performance Indicator
  15. Evaluating $\log p^\theta_{(K)}(y)$
  16. ...and 25 more sections

Key Result

Lemma 4.1

Suppose there exists $\tau^* > 0$ such that for all $x \in \mathbb{R}^d$, $0 \leq s < t \leq 1$, and $\tau \leq \min\{ \tau^*, t-s \}$. Then, $\phi^{\theta}_{s,t}(x) = \phi^{\eta}_{s,t}(x)$ and so satisfies Lagrangian PDE.

Figures (9)

  • Figure 1: 2D datasets. Each has 1M points.
  • Figure 2: The effect of how the intermediate time $u$ is chosen on the evolution of the KL divergence between distributions of student models and teacher models during training. For the strategy that samples from $[s,t]$, we sample from $[s,t-\tau]$ instead to avoid the denominator in the TVM loss being too small.
  • Figure 3: The effect of the hyperparameter $\tau$ on the evolution of the KL divergence between distributions of student models and teacher models during training.
  • Figure 4: Architecture of the teacher flow matching model for 2D datasets.
  • Figure 5: Architecture of the student TTFM for 2D datasets.
  • ...and 4 more figures

Theorems & Definitions (10)

  • Lemma 4.1
  • Lemma 4.2
  • Theorem 4.3: informal
  • Lemma 1.1
  • proof
  • Lemma 1.2
  • proof
  • Theorem 1.3
  • proof
  • Theorem 2.1: instantaneous change of variable formula