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

2025-05-02

AI Summary

This work addresses the poor sample quality of flow matching models under few-step generation. We propose the Two-Time Flow Matching (TTFM) distillation framework. Its core innovation is the Initial/Terminal Velocity Matching (ITVM) loss, which decouples velocity estimation at the start and end times, enabling independent optimization of both endpoints. To stabilize target velocity learning, we incorporate Exponential Moving Average (EMA). This design overcomes key limitations of prior methods like LFMD—namely, their reliance on higher-order derivatives and coupled velocity estimation. TTFM supports single-step, evaluable sample transformation over arbitrary time intervals and is compatible with diverse datasets and model architectures. Extensive experiments on CIFAR-10 and CelebA-HQ demonstrate that TTFM achieves superior FID and LPIPS scores using significantly fewer sampling steps, consistently outperforming existing distillation baselines.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Time-Series/Data Streams
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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 emph{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.

Problem

• Distilling flow models by matching initial and terminal velocities
• Improving few-step generation performance in flow models
• Extending LFMD loss with ITVM for better model distillation

Innovation

• Two-timed flow model for sample transformation
• Initial/terminal velocity matching loss function
• Exponential moving averaging for model stabilization