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.