Research map / Image generation / Fast diffusion sampling
Beyond Trajectory Matching: Reflow with Marginal Distribution Alignment
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Adds a marginal-alignment KL regularizer (computed via augmented ODE log-density tracking and frozen teacher scores) to PeRFlow/reflow distillation, improving 4-step text-to-image generation on SD 1.5 and SDXL.
Problem
Trajectory-matching objectives in reflow-style distillation may not determine the endpoint marginal distributions induced by the student, so few-step generation quality can differ even at equal trajectory loss.
Approach
Augment PeRFlow's trajectory matching loss with a window-wise KL marginal-alignment term between the student-pushed noising marginal and the teacher marginal, computed via log-density change along the student ODE and stop-gradient teacher scores, with no auxiliary networks; a telescoping TV bound links local alignment to final-time discrepancy.
Key results
- On SD 1.5 COCO-2014, 4-step FID improves from 12.01 (PeRFlow) to 10.66 (MA-Reflow).
- On SD 1.5 COCO-2017, FID drops from 23.81 to 22.76 and CLIP rises from 30.24 to 31.03.
- On SDXL COCO-2014-10K, 4-step FID is 19.17 vs 20.99 for PeRFlow; on COCO-2017 it is 25.94 vs 27.06 (Rectified Diffusion 25.81 is best).
Surprising findings
- Continuing PeRFlow training from its converged checkpoint degrades FID on CelebA-HQ (18.17 to 23.87 at 4 steps), while MA-Reflow continues to improve.
Limitations
- Evaluation of other reflow settings (different noise schedules, K=1) left to future work due to compute.
- Not applicable as-is to non-ODE student models such as consistency or shortcut models.
- Gains on SDXL COCO-2017 are not best (Rectified Diffusion slightly better); improvements are modest in some settings.
Experimental setup
Datasets: LAION-Aesthetics v2 5+ MS-COCO 2014 validation MS-COCO 2017 validation CelebA-HQ
Benchmarks: MS-COCO 2014 validation MS-COCO 2017 validation COCO-2014-10K CelebA-HQ
Metrics: FID CLIP score
Baselines: PeRFlow 2-ReFlow Flash Diffusion PCM LCM-LoRA LCM SDXL-Lightning SDXL-Turbo HyperSD Rectified Diffusion DMD2 DPMSolver
Builds on: Rectified Flow PeRFlow JKO-iFlow DDIM
Closest papers by meaning
- Optimal Flow Matching: Learning Straight Trajectories in Just One Step (NeurIPS 2024 · similarity 0.88)
- TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment (NeurIPS 2026 · similarity 0.88)
- Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation (NeurIPS 2026 · similarity 0.87)
- Shortcutting Pre-trained Flow Matching Diffusion Models is Almost Free Lunch (NeurIPS 2025 · similarity 0.86)
- Elucidating the Preconditioning in Consistency Distillation (ICLR 2025 · similarity 0.86)
- Learning to Solve Generative ODEs Beyond the Linear Span (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.