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

A flow matching model learns a time-dependent vector field vt(x) that generates a probability path pt 0 ≤ t ≤ 1 that interpolates between a well-known noise distribution (p0) and the data distribution (p1). 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 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.