这篇论文发现了再流蒸馏的隐藏问题——轨迹匹配可能不够，还提出了一个简单有效的边际对齐正则化，不用额外网络就能提升少步生成质量，值得做扩散加速的人看看。
现有扩散模型和连续流生成模型的确定性采样可视为求解学习到的ODE动力学，但准确离散化通常需要多步。论文指出轨迹匹配范式存在理论局限：两个学生模型可达到相同轨迹匹配损失却诱导不同端点边际分布，影响生成质量。为克服这一局限，提出边际对齐正则化器，通过追踪学生模型ODE沿线的对数密度变化并利用冻结教师模型评分来惩罚学生与教师边际分布的差异，无需辅助网络或对抗优化。该框架统一适用于原始再流和分段再流等再流族，并证明局部边际对齐通过 telescoping 总变差界控制最终时刻分布差异。在基准骨干网络上的实验验证了该方法在少步生成中的有效性。
Beyond Trajectory Matching: Reflow with Marginal Distribution Alignment
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.