# Beyond Trajectory Matching: Reflow with Marginal Distribution Alignment

arXiv: https://arxiv.org/abs/2606.29287 · HTML: https://mlsanity.com/paper/2606.29287/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 45 of its 48 references that are in the map.
- First posted on arXiv 28 Jun 2026.

## Builds on (most cited first, 30 of 45)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [InstructPix2Pix: Learning to Follow Image Editing Instructions](https://mlsanity.com/paper/2211.09800/) (2022)
- [Video Diffusion Models](https://mlsanity.com/paper/2204.03458/) (2022)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [SDEdit: Guided Image Synthesis and Editing with Stochastic Differential Equations](https://mlsanity.com/paper/2108.01073/) (2021)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Align your Latents: High-Resolution Video Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2304.08818/) (2023)
- [ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation](https://mlsanity.com/paper/2305.16213/) (2023)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [One-step Diffusion with Distribution Matching Distillation](https://mlsanity.com/paper/2311.18828/) (2023)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models](https://mlsanity.com/paper/2211.01095/) (2022)
- [Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference](https://mlsanity.com/paper/2310.04378/) (2023)
- [Adversarial Diffusion Distillation](https://mlsanity.com/paper/2311.17042/) (2023)
- [Improved Distribution Matching Distillation for Fast Image Synthesis](https://mlsanity.com/paper/2405.14867/) (2024)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation](https://mlsanity.com/paper/2309.06380/) (2023)

## Abstract (arXiv)

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.