The sampling procedures of diffusion and flow-matching models involve hundreds of denoising steps incurring substantial inference time, which has motivated a growing body of work on training schemes for few-step and one-step generation. Existing approaches primarily seek to align the generation trajectory of few-step and base models, yet their training objectives only involve discrepancies at a single point in the trajectory, encouraging only local point-level alignment. While recent methods such as Shortcut-CSL introduce cumulative objective that enforce alignment over the entire trajectory, this is suboptimal as they do not directly enforce alignment between the distributions induced by the two models throughout the generation process. We propose Cumulative Marginal Matching (CMM), a principled objective for few-step generation: by viewing flow-matching as a sequence of intermediate marginal distributions connecting the prior and data distributions, CMM directly optimizes the cumulative discrepancy between the base and few-step marginal distributions over the entire sequence. Specifically, we formulate this discrepancy using the Wasserstein-2 distance. By exploiting the pushforward representation of flow-matching, we transform an intractable cumulative Wasserstein objective into a tractable surrogate that can be optimized using samples from the intermediate marginal, which are readily generated through the conditional interpolation process of flow-matching. Furthermore, we show that the existing Shortcut-CSL objective arises as a first-order approximation of that of the CMM, establishing a theoretical connection between the two objectives. Experiments on standard benchmarks demonstrate improvement in generation quality over Shortcut-CSL and other state-of-the-art one- and few-step baselines under identical training budgets.
Figures & tables
Figure 1: Samples generated from CMM model trained on CelebA-256 with 128, two and one-step sampling. The two- and one-step preserve the sample quality of the base 128-step model.
Figure 2: Illustration of the CSL (left) and CMM (right) losses (the illustration of CSL is recreated based on the original Shortcut-CSL paper). In the CSL loss, few-step model is unfolded in time and its misalignment ( uθ ) with the model velocity field is penalized at each point. In CMM, we let the base and few-step model unfold together on their own trajectories, and penalize the difference between their trajectories ( ΔX ) at each point.
CelebA-256
ImageNet-256
Method
Four-Step
Two-Step
One-Step
Four-Step
Two-Step
One-Step
Reflow
13.77 ± 0.05
14.48 ± 0.05
16.07 ± 0.02
40.86
43.68
50.40
Progressive Distillation
14.49 ± 0.07
16.73 ± 0.02
20.40 ± 0.15
-
-
-
Flow-Matching
62.8 ± 0.04
112.1 ± 0.07
321.2 ± 0.04
104.96
210.39
325.78
Consistency Training
19.0
-
33.2
-
-
-
MeanFlow
-
-
-
34.09
36.61
45.12
Table 1: Comparison of the generation quality of the baselines and our method for unconditional generation on CelebA-256, and class-conditional generation on Imagenet-256 in terms of the FID-50K ↓ scores. The ratios after the methods represent the proportion of the number of data points used for the flow-matching and distillation parts of the objectives.
Training Time (in hours)
Method
DiT-S-2
DiT-B-2
FM
–
6.4
Shortcut
5.6
7.4
ST-CSL
6.2
7.8
CMM
6.6
8.6
Table 2: Comparison of the wall-clock training time of Flow Matching (FM), Shortcut , Shortcut-CSL (ST-CSL), and CMM over 100 training epochs on the CelebA-256 dataset using Diffusion Transformers of different sizes.
Hyperparam
CelebA-256
Imagenet-256
Dataset Size
∼ 30k
∼ 40k
#Epochs
300
500
Optimizer
AdamW
AdamW
EMA
0.9999
0.9999
B
64
128
K
16
32
Table 3: Hyperparameter setting for experiments in Tables 1 & 2 in the main paper.
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.
Chen Wang, Peiran Yun, Pan Xie +1
Department of Statistics and Data Science, Tsinghua University · Department of Computer Science and Technology, Tsinghua University · Bytedance Inc
We propose Perceptual Flow Matching (PFM), a simple yet effective framework for few-step generation in flow-matching models. Rather than performing velocity regression in the conventional VAE latent space, PFM supervises flow matching in a perceptual feature space using pretrained perceptual models. This simple change substantially improves the few-step generation capability of flow-matching models, reducing the number of sampling steps from 35-50 to 4-8 while preserving generation quality. Unlike existing acceleration and distillation approaches, PFM requires neither teacher models nor auxiliary score networks and can be integrated into standard flow-matching training pipelines with minimal modifications. Extensive experiments on image generation, video generation, and image editing tasks demonstrate that PFM consistently produces high-quality results while producing fewer artifacts than existing distillation-based methods. We further show that perceptual supervision shifts the regression minimizer from mean-seeking to mode-seeking, biasing predictions toward on-manifold modes that remain accurate under coarse few-step integration. Our results reveal that standard flow-matching training can naturally yield high-quality few-step generators when supervised in an appropriate representation space. We hope this insight inspires future research into representation-aware objectives for efficient generative modeling.
Chuyang Zhao, Yifei Song, Hongfa Wang +7
1Joy Future Academy · 2Fudan University · 3Tsinghua University +1
Flow Matching (FM) is a simulation-free method for learning a continuous, invertible flow that interpolates between two distributions, and in particular generates data from noise. Inspired by the variational nature of the diffusion process as a gradient flow, we introduce a stepwise FM model, Local Flow Matching (LFM), which sequentially learns a sequence of FM submodels, each matching a diffusion process up to the time-step size in the data-to-noise direction. In each step, the two distributions to be interpolated by the sub-flow model are closer than those in the full-flow matching model, which interpolates data to noise distributions, enabling smaller models with more efficient training. This variational perspective also allows us to prove a theoretical generation guarantee for the proposed flow model in terms of the χ2-divergence between the generated and true data distributions, leveraging the contraction property of the diffusion process. In practice, the stepwise structure of LFM is naturally amenable to model distillation, and various distillation techniques can be applied to accelerate generation. We empirically demonstrate that LFM achieves competitive generative performance compared to FM on unconditional generation of tabular and image datasets, and on conditional generation of robotic manipulation policies.
Chen Xu, Xiuyuan Cheng, Yao Xie
H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology. · Department of Mathematics, Duke University