Beyond Trajectory Matching: Reflow with Marginal Distribution Alignment

• URL: http://arxiv.org/abs/2606.29287v1
• Date: Sun, 28 Jun 2026 09:19:44 GMT
• Title: Beyond Trajectory Matching: Reflow with Marginal Distribution Alignment
• Abstract summary: Diffusion and continuous-flow generative models achieve high-quality generation.<n>Reflow-based distillation simplifies teacher ODE trajectories so that a student model can approximate the teacher transport with fewer steps.<n>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.
• Score: 7.13416637929041
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Distillation of Foundation Models for Time-dependent PDEs [21.455872116756492]
  We propose a knowledge distillation framework that transfers the predictive capability of a pretrained foundation model into a compact and efficient student.<n>Starting from a fine-tuned teacher, TREX augments limited downstream data by generating long synthetic trajectories through teacher rollouts.<n>The resulting students can match or surpass the teacher's accuracy while reducing the number of parameters by several orders of magnitude.
  arXiv  Detail & Related papers  (2026-08-12T11:24:08Z)
• LC-GRPO: Bridging Train-Inference Gap for Flow-Based GRPO with Langevin Correction [59.32834240526981]
  Flow-based generative models are typically sampled by solving a deterministic ordinary differential equation (ODE)<n>Existing GRPO methods for flow models replace the Euler-time ODE with aotropic differential equation (SDE) during training.<n>We introduce LC-GRPO, a flow-based GRPO framework with Langevin correction targeting the marginal distribution at the resulting timestep.
  arXiv  Detail & Related papers  (2026-08-06T04:47:57Z)
• Global Convergence of Gradient Descent for Score Matching in Gaussian Mixtures via Reverse Fisher Divergence [67.12978375116599]
  We study an alternative objective: the reverse Fisher divergence, where the expectation is taken with respect to the student distribution.<n>We prove global convergence guarantees under a $widetilde(1)$-separation assumption on the target means.<n>Our rely on a new Lyapunov-based analysis of the gradient descent dynamics, showing that the reverse Fisher divergence has a much more favorable optimization landscape than the forward Fisher divergence.
  arXiv  Detail & Related papers  (2026-06-18T07:34:33Z)
• One-Step Distillation of Discrete Diffusion Image Generators via Fixed-Point Iteration [20.457786704615454]
  Fixed-Point Distillation (FPD) is an end-to-end framework that constructs local correction targets by partially corrupting the student's one-step draft and refining it with a single teacher step.<n>FPD achieves competitive visual fidelity and structural alignment within a single inference step, narrowing the gap to multi-step teachers.
  arXiv  Detail & Related papers  (2026-05-20T17:59:10Z)
• DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models [55.01951088768769]
  DiffusionOPD is a new multi-task training paradigm for diffusion models based on Online Policy Distillation (OPD)<n>We show that DiffusionOPD consistently surpasses both multi-reward RL and cascade RL baselines in training efficiency and final performance.
  arXiv  Detail & Related papers  (2026-05-14T16:49:09Z)
• Don't Ignore the Tail: Decoupling top-K Probabilities for Efficient Language Model Distillation [50.19746127327559]
  We propose a new tail-aware divergence that decouples the contribution of the teacher model's top-K predicted probabilities from that of lower-probability predictions.<n> Experimental results demonstrate that our modified distillation method yields competitive performance in both pre-training and supervised distillation of decoder models.
  arXiv  Detail & Related papers  (2026-02-24T11:54:06Z)
• FlowSteer: Guiding Few-Step Image Synthesis with Authentic Trajectories [82.90132015584359]
  ReFlow has theoretical consistency with flow matching but suboptimal performance in practical scenarios.<n>We propose FlowSteer, a method unlocks the potential of ReFlow-based distillation by guiding the student along teacher's authentic generation trajectories.
  arXiv  Detail & Related papers  (2025-11-24T07:13:23Z)
• Knowledge Distillation of Uncertainty using Deep Latent Factor Model [10.148306002388196]
  We introduce a new method of distribution distillation called Gaussian distillation.<n>It estimates the distribution of a teacher ensemble through a special Gaussian process called the deep latent factor model (DLF)<n>By using multiple benchmark datasets, we demonstrate that the proposed Gaussian distillation outperforms existing baselines.
  arXiv  Detail & Related papers  (2025-10-22T06:46:59Z)
• Toward Theoretical Insights into Diffusion Trajectory Distillation via Operator Merging [10.315743300140966]
  Diffusion trajectory distillation aims to accelerate sampling in diffusion models that produce high-quality outputs but suffer from slow sampling speeds.<n>We propose a programming algorithm to compute the optimal merging strategy that maximally preserves signal fidelity.<n>Our findings enhance the theoretical understanding of diffusion trajectory distillation and offer practical insights for improving distillation strategies.
  arXiv  Detail & Related papers  (2025-05-21T21:13:02Z)
• Elucidating the Preconditioning in Consistency Distillation [25.213664260896103]
  We propose a principled way dubbed textitAnalytic-Precond to analytically optimize the preconditioning according to the consistency gap.<n>We demonstrate that Analytic-Precond can facilitate the learning of trajectory jumpers, enhance the alignment of the student trajectory with the teacher's, and achieve $2times$ to $3times$ training acceleration of consistency trajectory models.
  arXiv  Detail & Related papers  (2025-02-05T06:30:37Z)
• Multi-Granularity Semantic Revision for Large Language Model Distillation [66.03746866578274]
  We propose a multi-granularity semantic revision method for LLM distillation. At the sequence level, we propose a sequence correction and re-generation strategy. At the token level, we design a distribution adaptive clipping Kullback-Leibler loss as the distillation objective function. At the span level, we leverage the span priors of a sequence to compute the probability correlations within spans, and constrain the teacher and student's probability correlations to be consistent.
  arXiv  Detail & Related papers  (2024-07-14T03:51:49Z)
• Distributionally Robust Models with Parametric Likelihood Ratios [123.05074253513935]
  Three simple ideas allow us to train models with DRO using a broader class of parametric likelihood ratios. We find that models trained with the resulting parametric adversaries are consistently more robust to subpopulation shifts when compared to other DRO approaches.
  arXiv  Detail & Related papers  (2022-04-13T12:43:12Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.