Drift Flow Matching

Abstract: Iterative generative models such as Flow Matching and Diffusion models have demonstrated strong test-time scaling behavior, where additional inference computation can improve generation quality. In contrast, Drift Models offer efficient one-step generation, but their direct generation paradigm limits such flexibility. In this work, we propose Drift Flow Matching (DFM), a framework that connects drifting generative modeling with flow-based iterative generation. DFM preserves the efficiency of direct transport maps while enabling generation to be refined through multiple inference steps when desired. This bridges the gap between one-step Drift Models and multi-step Flow Matching methods, and provides a novel generative paradigm that can adapt sampling computation to different quality--efficiency requirements. Extensive experiments across different tasks and datasets demonstrate the effectiveness and generality of the proposed framework.

• The paper introduces a unified generative modeling framework that bridges efficient one-step drift models with iterative flow matching.
• It employs path construction with marginal transport via a mean-velocity field and drift-based supervision for robust distribution alignment.
• Empirical results on synthetic benchmarks, image synthesis, and robotic control highlight DFM’s tunability between speed and quality.

Drift Flow Matching: Unifying One-Step Drift Models with Iterative Flow Matching

Overview and Motivation

Iterative generative modeling frameworks such as Flow Matching (FM) and Diffusion Models leverage simulation of continuous-time dynamics at inference, enabling the quality of sample generation to improve as additional computation is spent. However, Drift Models pursue direct distribution transport via learned pushforward maps, optimizing for efficient one-step generation. These paradigms previously stood in contrast: Drift Models prioritized efficiency without inference-time refinement; Flow Matching and Diffusion models enabled flexible scaling but at considerable computational expense. "Drift Flow Matching" (DFM) (2605.17244) introduces a theoretical and algorithmic nexus that unites these two approaches, establishing a flexible generative framework that interpolates between fast (one-step) and high-quality (multi-step) generation, governed by a learned distribution-level drift objective on marginal pairs.

Methodological Foundations

Path Construction and Marginal Transport

DFM adopts the conditional path machinery from FM, constructing interpolants between source and target distributions indexed by endpoint pairings. For any two time-points $(t, r)$ along the generative trajectory (with $0 < t < r < 1$), the framework leverages the same endpoint-independent coupling as FM and forms marginal distributions $p_t$ and $p_r$ by pathwise interpolation. The transport map is parameterized by a mean-velocity field $u_\theta(x_t, t, r)$, enabling direct movement from $p_t$ to $p_r$ in a single evaluation:

$$T_{t, r}^\theta(x_t) = x_t + (r - t) u_\theta(x_t, t, r)$$

This construction recovers the FM dynamics in the infinitesimal-step limit, guaranteeing consistency with classical probability flow ODEs.

Drift-Based Supervision

Distinct from FM's velocity regression, DFM introduces drift supervision at the level of marginal distributions. The drift field $V_{q_{t, r}, p_r}$ is computed using kernel-weighted interactions between predicted samples (from the pushforward of $p_t$) and true target samples from $0 < t < r < 1$0, with both attraction to the target and self-correction effects. The stop-gradient drift objective for a time-pair $0 < t < r < 1$1 is:

$0 < t < r < 1$2

This formulation enables learning large-step transports directly, circumventing the need for velocity field regression and offering robust distribution-level alignment.

Grouped Drift Estimation

DFM requires estimations of the drift field to be performed group-wise for each specific $0 < t < r < 1$3 pair, as mixing samples across time-pairs invalidates the marginal transport structure. Mini-batch training is structurally compatible, and ablations show that a small number of time pairs per batch suffices for efficient optimization.

Practical Results and Numerical Highlights

Synthetic and Conditional Image Generation

DFM achieves strong one-step sample quality on classic synthetic benchmarks, producing accurate transport to target distributions (e.g., "F", "M", two-moons, checkerboard shapes) while enabling further improvements in distribution coverage and fidelity as the number of inference steps (NFE) increases.

In class-conditional MNIST and FFHQ generation, DFM matches Drift Model baselines at NFE=1 and demonstrates test-time scaling: increasing NFE yields lower 2-Wasserstein distances and higher generation accuracy, validating the interpolation between Drift and FM paradigms. On FFHQ, both PCA and UMAP latent space visualizations show improved class separation and coverage as NFE grows.

Scalable Image Synthesis (ImageNet)

DFM is competitive with state-of-the-art models on high-resolution ImageNet-1k synthesis when compared at equal compute. With NFE=1, DFM matches Drift Model results in terms of FID and IS. As NFE increases, sample quality improves, approaching the performance of multi-step FM and diffusion transformer models. This provides a unique trade-off: rapid sampling with optional refinement.

Robotic Control

Replacing the policy core in Diffusion Policy and Drift Policy frameworks with DFM yields high success rates across both single-stage and multi-stage robotic manipulation benchmarks. DFM preserves one-step inference speed and introduces test-time scaling capability, confirming its versatility beyond image domains.

Ablation Analyses

• Time Embedding: Explicit encoding of current time and step size yields the best results; direct embedding of $0 < t < r < 1$4 provides stable model conditioning for transport.
• Time Pair Sampler: Logit-normal distributions outperform uniform, echoing findings in FM literature.
• Kernel Temperature: Intermediate temperature values (e.g., $0 < t < r < 1$5) optimize drift supervision; too low/high temperatures degrade sample specificity or introduce noise sensitivity.
• Positive/Negative Samples: Grouping sufficient positive/negative samples per time-pair is critical for reliable drift estimation and performance in small NFE regimes.
• Model Parameterization: Mean-velocity parameterization outperforms direct target-state prediction, facilitating compatibility with both large and infinitesimal steps.

Theoretical Implications

DFM inherits explicit $0 < t < r < 1$6-geometric structure from the FM path construction, with provable upper bounds on Wasserstein transport cost between marginals for arbitrary $0 < t < r < 1$7. The discrete action along the DFM path decomposes the endpoint transport into controlled short-range subproblems, avoiding excess quadratic cost. The framework admits a first-order expansion in the infinitesimal-step limit, rigorously recovering the FM velocity field. This theoretical foundation assures both local and global consistency with established generative modeling principles and optimal transport theory.

Implications and Future Directions

DFM establishes a new generative modeling paradigm where efficiency and quality are not dichotomous but are tunable per application. For real-world deployment, DFM allows inference-time adaptation: practitioners can select generation speed (one-step) or choose iterative refinement (multi-step), targeting quality or efficiency demands dynamically. This is especially impactful in conditional settings, high-dimensional structured data (e.g., robotic control, protein design), and domains where compute resources vary.

Theoretically, DFM provides a unified lens to analyze drift-based generative models and their relation to FM and diffusion methods. The group-wise distribution-level drift supervision could further be extended to optimal transport and Sinkhorn-based variants for maximized identifiability and sample coherence. As large-scale generative models expand to text, audio, and control domains, the DFM principle will likely inform algorithmic design for flexible, scalable inference.

Current limitations stem from inherited issues of generative modeling: exposure to data bias, privacy, and systematic error. For safety-critical deployment (e.g., synthetic media generation or robotics), robust data governance and oversight will be necessary.

Conclusion

"Drift Flow Matching" introduces a theoretically grounded, practically flexible generative modeling method that bridges efficient single-step drift models and scalable iterative flow matching. It expands the landscape for generative modeling, enabling controlled trade-offs between computational efficiency and sample quality. Empirical and theoretical results confirm DFM’s generality and robustness across image synthesis and control, setting the stage for future research into adaptive generative paradigms.