PaperlayerView exact revision on arXiv

Source-linked AI summary

MixFlow: Mixed Source Distributions Improve Rectified Flows

Nazir Nayal, Christopher Wewer, Jan Eric Lenssen

arXiv:2604.09181v1cs.CVcs.LG

TL;DR

Slow iterative sampling in flow models is linked to curved trajectories arising from weak source–data alignment. The paper introduces κ-FC and MixFlow to improve that alignment, reporting 12% and 7% FID improvements over specified baselines with comparable sampling steps.

  • Problem

    Flow models use source distributions that may be poorly aligned with data, while curved generative paths make high-quality sampling slow.

  • Method

    κ-FC learns source distributions conditioned on arbitrary signals, while MixFlow trains on mixtures of fixed unconditional and κ-FC-based distributions.

  • Results

    12% FID improvement over standard Rectified Flow and 7% over the best previous trajectory-straightening method are reported with comparable sampling steps.

  • Takeaways & Limitations

    MixFlow reduces trajectory curvature, improves performance under a fixed sampling budget, and accelerates training convergence.

  • Takeaways & Limitations

    MixFlow still requires a Gaussian assumption, and the authors identify relaxing this assumption as future work.

Abstract

from arXiv · show

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 $κ\texttt{-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 $κ\texttt{-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: $\href{https://github.com/NazirNayal8/MixFlow}{https://github.com/NazirNayal8/MixFlow}$

1 INTRODUCTION

The paper targets slow sampling in flow models by improving source–data alignment and reducing generative-path curvature. It introduces κ-FC and MixFlow, reporting better FID with comparable sampling steps.

  • Motivation: Flow models can require many forward passes because sampling quality depends on straightening their generative paths.Rectified Flow and Flow Matching address this by using optimal-transport displacement interpolations.
  • Contributions: κ-FC is a learnable forward-coupling formulation that uses arbitrary conditioning signals to align the source distribution with the target.The signal can provide information related to the data distribution.
  • Contributions: MixFlow trains flows with a linear mixture of a fixed unconditional distribution and a κ-FC-based learnable distribution.The mixture transports trajectory structure from the conditional source to the unconditional Gaussian source.
  • Results: 12% FID improvement over standard Rectified Flow and 7% over the best previous trajectory-straightening method are reported with comparable sampling steps.The evaluation uses common image-generation benchmarks.
  • Contributions: The contributions are κ-FC for better source distributions and MixFlow for requiring fewer steps to generate high-quality samples.These are listed as the paper’s two main contributions.

2 RELATED WORK

Related work improves sampling through distillation, faster solvers, path straightness, or optimized forward couplings. MixFlow instead optimizes source–target dependence while remaining compatible with faster-solver approaches.

  • Distillation: Distillation methods linearize source-to-target mappings through distillation, consistency constraints, or Reflow.These approaches seek reasonable generation quality with fewer sampling steps.
  • Faster Solvers: Faster-solver methods improve numerical ODE integration, but sampling speed remains bounded by curvature induced by the learned trajectories.The source-distribution approach is presented as orthogonal and therefore combinable with faster solvers.
  • Path Straightness: Rectified Flow and Flow Matching straighten trajectories through path-intersection analysis and optimal-transport interpolation.Other path-straightness methods also target trajectory curvature.
  • Path Straightness: Prior path-straightening methods still assume an independent source–target coupling, whereas this work reduces curvature by optimizing that coupling.The paper frames source-distribution optimization as distinct from these approaches.
  • Optimized Forward Coupling: Optimized-forward-coupling methods approximate optimal-transport plans between source and target distributions, often at the mini-batch level.The approximation is motivated by the computational infeasibility of solving optimal transport over an entire dataset.

3 BACKGROUND

Rectified Flow learns a vector field that maps a tractable source distribution to a complex target along straight interpolations. Forward-coupling choices affect path intersections, curvature, and sampling speed.

  • Rectified Flow: Rectified Flow maps samples from a tractable source distribution p0(x) to a complex target distribution p1(x).Their joint coupling preserves the respective marginal densities.
  • Rectified Flow: The intermediate state is xt = tx1 + (1−t)x0, and its required velocity is the constant derivative x1 − x0.This defines a straight path for t ∈ [0, 1].
  • Rectified Flow: The learned vector field vθ(xt, t) approximates the velocity needed to transport source samples to target samples along these paths.Its parameters are learned by minimizing the Rectified Flow objective.
  • Trajectory Curvature: Highly intersecting forward paths make the vector field estimate mean directions, producing highly curved generative trajectories.The coupling therefore influences the curvature learned by the model.
  • Trajectory Curvature: With a fixed independent coupling, the intersection degree cannot change, so prior methods optimize q(x0, x1) to reduce intersections and improve sampling speed.The optimal vector field is the conditional mean direction given the intermediate state.

4 STRAIGHTENED TRAJECTORIES VIA DISTRIBUTION MIXING

This section introduces κ-FC for learnable source distributions conditioned on arbitrary signals, then uses distribution mixing to improve stability, curvature, and sampling efficiency. MixFlow combines a κ-conditioned Gaussian with a standard Gaussian during training and supports inference-time control through the mixing weight.

  • κ-FC: κ-FC models a learnable forward coupling in which an arbitrary signal κ conditions the source distribution to align it with the data.κ may be informative, such as a class label or image caption, or independent; the formulation includes data-sample conditioning as a special case.
  • κ-FC: A naively parameterized conditional Gaussian requires KL regularization whose weight β creates a trade-off between prior coverage and κ-dependent alignment.As β approaches zero, the prior hole problem prevents inference without κ; as β approaches infinity, the source becomes nearly standard Gaussian and loses its dependence on κ.
  • Distribution mixing: MixFlow trains on Gaussian interpolations between N(µϕ(κ), Σϕ(κ)) and N(0, I), producing qϕ(x | κ, w) = N(wµϕ(κ), wΣϕ(κ) + (1 − w)I).The mixing weight w is sampled independently for each training example, so the vector field learns flows from conditional, unconditional, and intermediate source distributions.
  • Results: ∼5% improved curvature over Fast-ODE is reported for MixFlow with a lower KL-regularization coefficient β.Table 1 compares MixFlow with Rectified Flow and Fast-ODE using generative trajectory curvature.
  • Distribution mixing: Mixing permits β values on the order of 10^-5 while retaining training stability and Gaussian-prior coverage, allowing greater source-distribution complexity and lower trajectory curvature.The KL term is R(x0, κ, w) = DKL(qϕ(x0 | κ, w)∥N(0, I)).
  • Sampling: At inference, κ-conditioned Gaussian interpolants or the standard normal can initialize the ODE, while w controls the speed-quality trade-off without retraining.The standard normal serves as a fallback when κ is unavailable.

5 EXPERIMENTS

Experiments evaluate MixFlow on CIFAR10, FFHQ, and AFHQv2 across curvature, sampling-budget, solver, and resolution settings. MixFlow improves FID and trajectory curvature while offering a strong speed-quality trade-off without retraining across different regularization choices.

  • 5.1.1 CURVATURE EVALUATION: MixFlow achieves ∼22% lower trajectory curvature than Rectified Flow and ∼5% lower curvature than Fast-ODE.The curvature evaluation uses 10K trajectories and an Euler sampler with 128 inference steps.
  • 5.1.2 GENERATION EVALUATION: Table 2 evaluates CIFAR10 FID with RK45 full simulation and Heun’s solver at 5 and 9 function evaluations.MixFlow improves sample quality with small numbers of function evaluations.
  • 5.1.2 GENERATION EVALUATION: MixFlow reduces FID by ∼12% versus standard Rectified Flow and ∼7% versus Fast-ODE at comparable NFE.The full-simulation comparison evaluates CIFAR10 with ODE solvers including RK45.
  • 5.1.2 GENERATION EVALUATION: At NFE = 5, MixFlow improves FID by ∼20% versus Fast-ODE and ∼2% versus QAC.At NFE = 9, it improves FID by ∼10% versus Fast-ODE and ∼12.7% versus QAC.
  • 5.2 UNCONDITIONAL GENERATION ON FFHQ & AFHQ: On FFHQ and AFHQv2 at 64 × 64, MixFlow outperforms Fast-ODE for almost all NFEs and provides the best speed-quality trade-off without retraining for different β values.Qualitative generations on both datasets reach reasonable quality with fewer than 10 steps.

6 ANALYSIS

The analysis studies how regularization and conditioning signals shape MixFlow’s source distribution and generation quality. Lower β improves FID until source collapse occurs, while more informative κ generally improves alignment and performance.

  • 6.1 CHOICE OF CONDITIONING FOR SOURCE DISTRIBUTION: κ = x1 is the most informative conditioning choice because it uses the data sample itself.The analysis also considers class labels κc and independent Gaussian noise κn.
  • 6.1 CHOICE OF CONDITIONING FOR SOURCE DISTRIBUTION: All tested conditioning signals improve FID over Rectified Flow, including uninformative κn sampled from a standard Gaussian.The learnable forward coupling can map standard-Gaussian samples to a data-aligned sub-region.
  • 6.1 CHOICE OF CONDITIONING FOR SOURCE DISTRIBUTION: Lower β values improve FID across sampling steps, with clear improvements at β = 10−5.Regularization remains necessary because source distributions collapse for β values around 10−8.

7 CONCLUSION

The paper addresses flow-model sampling efficiency through curvature minimization, proposing κ-FC and MixFlow to improve fixed-budget performance and training convergence. It also identifies remaining Gaussian assumptions as a limitation for future work.

  • κ-FC provides a general learnable forward-coupling formulation that can leverage arbitrary signals in rectified flows.
  • MixFlow mixes conditional and unconditional distributions during training to minimize trajectory curvature, improve fixed-budget performance, and accelerate convergence.
  • MixFlow still requires a Gaussian assumption, motivating future work on further relaxations while maintaining performance.

Supplementary Materials

The supplementary analysis estimates generated-trajectory curvature from ODE-solved trajectories and examines how the KL regularization weight β affects curvature. Lower β values empirically correlate with lower curvature.

  • The curvature evaluation generates 10,000 trajectories with a 128-step Euler solver and averages curvature for an optimized vector-field model.
  • The source endpoint is sampled from p0(x) = N(0, I), while the target endpoint is obtained as x1 = ODESolver(x0, vθ).
  • The curvature measure differs from the coupling-based intersection degree because its vector field is optimized and its endpoint is computed deterministically by an ODE solver.
  • Lower KL regularization weights β empirically correlate with lower generative-trajectory curvature across MixFlow models.

B IMPLEMENTATION DETAILS

The implementation uses UNet-based vector-field and source-prediction networks, with configurations reported by dataset. Supplementary figures examine curvature as β changes and FID as the interpolation weight varies.

  • The vector-field model uses the Fast-ODE UNet architecture following the DDPM++ implementation of EDM.
  • The source-prediction network is a small UNet whose input handling changes for class-label conditioning through an appended embedding layer.
  • Figure 3 shows lower trajectory curvature with lower β values.
  • Table 6 reports vector-field UNet configurations and training hyperparameters separately for each dataset.
  • Table 7 reports the UNet configuration for the source-prediction parametrization, following the diffusers library definition.
  • Figure 4 varies interpolation weight w during sampling and compares FID across sampling-step counts for class-label and noise conditioning.

C.1 TRAINING EFFICIENCY

MixFlow improves training and sampling efficiency, reaching Fast-ODE-level performance with 60% of the training budget and improving low-step generation. The useful mixture weight depends on signal informativeness and sampling budget.

  • Training efficiency: 60% of the full training iterations suffices for MixFlow to achieve approximately the same performance as Fast-ODE.
  • Figure interpretation: The figures compare FID across sampling steps and mixture weights, with lower FID represented by stronger red and higher FID by stronger blue.
  • Effect of w: With uncorrelated Gaussian-noise conditioning, FID is best at w = 0 and degrades as w increases toward 1.
  • Effect of w: When w = 0, noise conditioning performs slightly better than class-label conditioning, but increasing w makes class-label conditioning better across sampling-step counts.
  • Qualitative generation: MixFlow generates higher-quality samples than Rectified Flow at low sampling steps, specifically 2 and 4.

D ADDITIONAL QUALITATIVE RESULTS

The section presents qualitative MixFlow generations across CIFAR10, FFHQ 64 × 64, and AFHQv2 64 × 64, including comparisons with Rectified Flow at different sampling budgets.

  • Qualitative results: MixFlow qualitative examples cover CIFAR10, FFHQ 64 × 64, and AFHQv2 64 × 64.The examples use κ as the data sample, β = 10−5, an Euler solver, and 64 sampling steps.
  • Comparison with Rectified Flow: At 2 and 4 sampling steps, MixFlow generates much clearer images than Rectified Flow.Figure 7 compares generated images from both methods across different sampling-step counts.
  • Qualitative results: The additional figures provide qualitative results for CIFAR10, FFHQ 64 × 64, and AFHQv2 64 × 64.These dataset-specific results are shown in Figures 8, 9, and 10, respectively.
Loading 2604.09181v1…