Diffusion models and their variations, such as rectified flows, generate diverse and high-quality images, but they are still hindered by slow iterative sampling caused by the highly curved generative paths they learn. An important cause of high curvature, as shown by previous work, is independence between the source distribution (standard Gaussian) and the data distribution. In this work, we tackle this limitation by two complementary contributions. First, we attempt to break away from the standard Gaussian assumption by introducing κ-FC, a general formulation that conditions the source distribution on an arbitrary signal κ that aligns it better with the data distribution. Then, we present MixFlow, a simple but effective training strategy that reduces the generative path curvatures and considerably improves sampling efficiency. MixFlow trains a flow model on linear mixtures of a fixed unconditional distribution and a κ-FC-based distribution. This simple mixture improves the alignment between the source and data, provides better generation quality with less required sampling steps, and accelerates the training convergence considerably. On average, our training procedure improves the generation quality by 12% in FID compared to standard rectified flow and 7% compared to previous baselines under a fixed sampling budget. Code available at: https://github.com/NazirNayal8/MixFlow
Figures & tables
Figure 1: Method overview. We propose training rectified flows with mixed source distributions, obtained by interpolating a conditional and a simple unconditional distribution. The conditional distribution is predicted from a signal κ , which is possibly informative, e.g. a specific data example, a class label, or entirely independent, e.g. random noise. The learned conditional distribution provides a trajectory structure that minimizes the degree of intersections, which is inherited by the mapping from the unconditional source distribution.
Model
β
Curvature ( ↓ )
Rectified Flow
∞
0.0467
Fast-ODE
20
0.0388
MixFlow (ours)
10−5
0.0366
Table 1: Trajectory curvature. We compare the generative trajectory of MixFlow with Rectified Flow and Fast-ODE. With lower β coefficient for the KL divergence loss, MixFlow achieves ∼5% improved curvature over Fast-ODE.
Method
Solver
NFE
FID ( ↓ )
Rectified Flow
RK45
127
2.58
FM - OT
142
6.36
Minibatch-OT
133.9
3.58
Fast-ODE
118
2.45
QAC
-
2.43
Ours
124.7
2.27
Table 2: Speed-quality trade-off on CIFAR10. We evaluate our method on CIFAR10 in terms of FID using different ODE solvers. In the top, we fully simulate the ODE trajectory using the RK45 adaptive solver. MixFlow achieves notable improvements in FID with a comparable number of function evaluations (NFEs). Our method further improves sample quality with a small number of NFE as can be seen for 5 and 9 function evaluations with Heun’s 2nd order solver.
Model / NFE
β
4
5
10
20
32
64
128
FFHQ
Fast-ODE
10
32.58
25.33
13.21
8.85
7.54
6.91
7.01
20
38.23
29.12
14.03
8.78
7.08
5.95
5.72
30
41.16
30.75
14.37
8.76
6.90
5.45
4.93
Ours
5×10−5
33.72
25.04
12.23
7.52
5.31
4.01
3.75
Table 3: Comparison with Fast-ODE. using with different KL divergence weights β w.r.t FID-10k on the FFHQ and AFHQv2 64×64 datasets. Our model outperforms Fast-ODE for different β choices on almost all NFEs. MixFlow provides overall the best trade-off between sampling speed and quality without the need to retrain with a different β parameter.
Figure 2: Qualitative Results. We show our method’s generation on FFHQ (rows 1-2) and AFHQv2 (rows 3-4) datasets using different steps. MixFlow requires few steps to generate reasonable outputs.
Table 6
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Curvature vs. β . We show how the curvature of the generative trajectories changes with different β values. We can see a clear trend of lower curvature with a lower β value
Parameters
CIFAR10
FFHQ
AFFHQv2
UNet Parameters
Channel Size
128
128
128
Channel Multiplier
[2,2,2]
[1,2,2,2]
[1,2,2,2]
Blocks per Layer
4
4
4
Attention Resolution
16
16
16
Dropout Probability
0.13
0.05
0.25
Embedding type
positional
positional
positional
Appendix
Table 6: Model and Experiments Configurations. In the upper part, we show the UNet configuration for the vector field model vθ . In the lower part we show the training hyperparameters. Each column show the configuration for a specific dataset.
Table 7: Source Prediction Network qϕ(x0∣κ) Configuration. We show the UNet configurations for the source prediction parametrization. The params follow the diffusers library definition of the UNet model.
Figure 4: Effect of interpolation weight w . We visualize the effect of varying w (x-axis) during sampling on FID (y-axis) across different numbers of sampling steps (different lines). (a) Sampling with a few steps benefits from a larger weight of the source distribution conditioned on class label κc , while for many steps, the unconditional standard Gaussian is better suited. (b) With conditioning on uncorrelated Gaussian noise κn , the best FID is achieved for w=0 , i.e., not using the conditional distribution at all during sampling. However, note that even in this case training the vector field on a mixture of distributions still improves performance as shown in Tab. 5 .
Figure 5: FID for Sampling Steps vs. weight w . We show the FID across different sampling step choices for both κc and κn as the interpolation parameter w changes. This is the numerical version of Figure 3 in the main paper, which provides a more accurate look into the change of FID values, where red indicates lower FID and blue indicates higher FID.
Figure 6: FID vs. training progress. Samples are generated with the RK45 solver across different step of the training process. Our method achieves the same performance as Fast-ODE (gray dotted line) with only 60% of the training budget.
Figure 7: Comparison against Rectified Flow . We show multiple examples for generated images with different sampling steps and compare against Rectified Flow. We highlight that for a low number of sampling steps (2,4), MixFlow generates much clearer images compared to Rectified Flow.
Rectified flow is a generative modeling method that learns near-straight transports from noise to data for efficient sampling and conditional generation.
Recently, flow-based generative models have shown superior efficiency compared to diffusion models. In this paper, we study rectified flow models, which constrain transport trajectories to be linear from the base distribution to the data distribution. This structural restriction greatly accelerates sampling, often enabling high-quality generation with a single Euler step. Under standard assumptions on the neural network classes used to parameterize the velocity field and data distribution, we prove that rectified flows achieve sample complexity O~(ε−2). This improves on the best known O(ε−4) bounds for flow matching model and matches the optimal rate for mean estimation. Our analysis exploits the particular structure of rectified flows: because the model is trained with a squared loss along linear paths, the associated hypothesis class admits a sharply controlled localized Rademacher complexity. This yields the improved, order-optimal sample complexity and provides a theoretical explanation for the strong empirical performance of rectified flow models.
Flow-based models learn a target distribution by modeling a marginal velocity field, defined as the average of sample-wise velocities connecting each sample from a simple prior to the target data. When sample-wise velocities conflict at the same intermediate state, however, this averaged velocity can misguide samples toward low-density regions, degrading generation quality. To address this issue, we propose the Flow Divergence Sampler (FDS), a training-free framework that refines intermediate states before each solver step. Our key finding reveals that the severity of this misguidance is quantified by the divergence of the marginal velocity field that is readily computable during inference with a well-optimized model. FDS exploits this signal to steer states toward less ambiguous regions. As a plug-and-play framework compatible with standard solvers and off-the-shelf flow backbones, FDS consistently improves fidelity across various generation tasks including text-to-image synthesis, and inverse problems.
This paper studies the training-testing discrepancy (a.k.a. exposure bias) problem for improving the diffusion models. During training, the input of a prediction network at one training timestep is the corresponding ground-truth noisy data that is an interpolation of the noise and the data, and during testing, the input is the generated noisy data. We present a novel training approach, named MixFlow, for improving the performance. Our approach is motivated by the Slow Flow phenomenon: the ground-truth interpolation that is the nearest to the generated noisy data at a given sampling timestep is observed to correspond to a higher-noise timestep (termed slowed timestep), i.e., the corresponding ground-truth timestep is slower than the sampling timestep. MixFlow leverages the interpolations at the slowed timesteps, named slowed interpolation mixture, for post-training the prediction network for each training timestep. Experiments over class-conditional image generation (including SiT, REPA, and RAE) and text-to-image generation validate the effectiveness of our approach. Our approach MixFlow over the RAE models achieve strong generation results on ImageNet: 1.43 FID (without guidance) and 1.10 (with guidance) at 256 x 256, and 1.55 FID (without guidance) and 1.10 (with guidance) at 512 x 512.
Hui Li, Fu-Yun Wang, Haoyuan Xia +4
Shanghai Innovation Institute · Fudan University · Shanghai Academy of AI for Science +2