MixFlow: Mixed Source Distributions Improve Rectified Flows
| Authors | Nazir Nayal et al. |
| Year | 2026 |
| HF Upvotes | 0 |
| arXiv | 2604.09181 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
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{https://github.com/NazirNayal8/MixFlow}
Engineering Breakdown
Plain English
This paper addresses a fundamental inefficiency in diffusion models and rectified flows: the sampling process requires many iterative steps because the generative paths learned by these models are highly curved. The authors identify that this curvature problem stems from a mismatch between the source distribution (typically a standard Gaussian) and the actual data distribution. They propose two solutions: κ-FC, a flexible framework that conditions the source distribution on an arbitrary signal to better align with the data, and MixFlow, a training strategy that directly reduces path curvature by mixing source distributions during training.
Core Technical Contribution
The paper makes two complementary technical innovations. First, κ-FC generalizes the source distribution selection beyond the standard Gaussian by conditioning it on an arbitrary alignment signal κ, enabling the source to be tailored to the specific data distribution characteristics. Second, MixFlow is a training methodology that mixes multiple source distributions during the learning process, empirically reducing the curvature of the learned generative paths and enabling faster sampling with fewer iterations. The key insight is that the independence assumption between source and data distributions creates unnecessarily curved trajectories, and breaking this independence through both distribution conditioning and mixture-based training yields straighter, more efficient paths from noise to data.
How It Works
The method operates as follows: instead of always sampling from a fixed standard Gaussian N(0, I), the κ-FC framework conditions the source distribution on a signal κ derived from or aligned with the data distribution, creating a source that is inherently closer to the data manifold. During training, MixFlow doesn't commit to a single source distribution but instead trains on a mixture of distributions, allowing the model to learn how to navigate paths from diverse starting points. This mixture-based training encourages the learned trajectory to be geometrically straighter because the model must find paths that work well across multiple source distributions rather than optimizing for a single standard Gaussian. The generative process then samples from one of these conditioned/mixed source distributions and follows the learned trajectory to the data distribution with reduced curvature. The conditioning signal κ could be derived from properties of individual data samples, making the source distribution adaptive to the specific data being generated.
Production Impact
For engineers building image generation systems, this approach directly translates to faster inference: reducing path curvature means the model can generate high-quality images in fewer sampling steps, cutting latency and compute cost proportionally. If a standard rectified flow requires 50-100 sampling steps, MixFlow could enable comparable quality in 20-30 steps, which is a 2-5x speedup on the critical sampling bottleneck. The integration is straightforward—this affects training procedures and the source distribution selection at inference time, not the underlying model architecture, so it can be retrofitted into existing diffusion pipelines. The trade-off is increased training complexity: the model must learn to handle multiple source distributions, which may require longer training or more careful hyperparameter tuning. For production systems, the latency savings are likely substantial enough to justify the engineering effort, especially for latency-sensitive applications like real-time image editing or interactive generation.
Limitations and When Not to Use This
The paper assumes that you have meaningful signals κ available to condition the source distribution—this requires prior knowledge about or ability to extract features from the data distribution, which may not always be straightforward for complex, high-dimensional data. The κ-FC framework's effectiveness is likely data-dependent; conditioning signals that work well for natural images may not transfer to medical imaging, 3D meshes, or other modalities. The paper does not clearly address what happens when the conditioning signal κ is poorly chosen or misaligned with the data—this could potentially make path curvature worse rather than better. Additionally, the abstract is incomplete, so the empirical results, the specific algorithms for selecting mixture components, and comparisons to other path-straightening approaches are not visible; the true computational overhead and theoretical guarantees of MixFlow are unclear from the given text.
Research Context
This work builds on prior research showing that curved generative paths in diffusion models cause slow sampling, and extends recent work on rectified flows (which aim for linear paths) by addressing the fundamental source-data distribution mismatch. It contributes to the broader effort to make diffusion-based generation faster and more practical, following work on ODE-based samplers, consistency models, and other trajectory optimization techniques. The paper opens a research direction around adaptive source distributions and mixture-based training strategies that could be applied beyond image generation to audio, video, 3D, and other modalities where path efficiency is critical. The idea of conditioning generative processes on data-aligned signals may also inform research in conditional generation, multimodal diffusion, and transfer learning in generative modeling.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::