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

TL;DR

This paper introduces a new loss function called ITVM for distilling flow models into two-timed flow models, improving the efficiency and quality of generative sampling in probability distribution interpolation tasks.

Contribution

The paper proposes the ITVM loss, extending previous methods by separately matching initial and terminal velocities, leading to better few-step generation performance.

Findings

Enhanced few-step generation accuracy

Improved model stability with EMA

Better performance across datasets and architectures

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) $\phi_{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.