Beyond Trajectory Matching: Reflow with Marginal Distribution Alignment

2026-06-28

AI Summary

Existing trajectory-matching-based Reflow distillation methods suffer from unstable few-step generation quality due to the non-uniqueness of the student model’s marginal distributions. This work proposes a marginal distribution alignment regularization that requires neither additional trainable networks nor adversarial optimization. By tracking the log-density evolution along the student’s ODE trajectory and leveraging the frozen teacher’s score function, the method explicitly enforces consistency between the student and teacher marginal distributions at the endpoints of each distillation interval. Theoretical analysis shows that such local marginal alignment effectively controls the total variation error of the final distribution. Experiments demonstrate that this approach significantly improves few-step generation quality on standard backbone architectures and exhibits broad applicability across Reflow-based frameworks.

Abstract

Diffusion and continuous-flow generative models achieve high-quality generation, and their deterministic sampling can be formulated as solving learned ODE dynamics. However, accurate ODE discretization often requires many steps, making efficient few-step generation a key challenge. Among acceleration strategies, reflow-based distillation simplifies teacher ODE trajectories so that a student model can approximate the teacher transport with fewer steps. We identify a theoretical limitation of this paradigm, namely that trajectory matching can under-determine the distribution induced by the student model. In particular, two student models can attain the same trajectory-matching loss while inducing different endpoint marginal distributions, which may lead to different generation quality. To address this limitation, we introduce a marginal-alignment regularizer that penalizes the discrepancy between the student-induced marginal and the corresponding teacher marginal at the endpoint of each distillation interval. The regularizer is computed by tracking log-density changes along the ODE induced by the student model and evaluating scores from the frozen teacher model, without requiring auxiliary trainable networks or adversarial optimization. The resulting framework applies uniformly to the reflow family, including vanilla reflow and piecewise reflow. We further prove a telescoping total-variation bound showing that local marginal alignment controls the final-time discrepancy between the student-induced and teacher-induced distributions. Experiments on benchmark backbones demonstrate the effectiveness of the proposed method for few-step generation.

Problem

• trajectory matching
• marginal distribution
• reflow distillation
• generative models
• ODE dynamics

Innovation

• reflow distillation
• marginal distribution alignment
• trajectory matching
• continuous-flow generative models
• ODE-based generation