💻 computer science

Drift Flow Matching

This paper proposes Drift Flow Matching (DFM), a novel framework that bridges the gap between efficient one-step Drift Models and flexible multi-step Flow Matching methods, enabling adaptive generation quality by allowing direct transport maps to be refined through iterative inference steps.

Original authors: Chenrui Ma, Xi Xiao, Lin Zhao, Tianyang Wang, Ferdinando Fioretto, Yanning Shen

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Chenrui Ma, Xi Xiao, Lin Zhao, Tianyang Wang, Ferdinando Fioretto, Yanning Shen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Core Problem: Speed vs. Quality

Imagine you are an artist trying to draw a perfect portrait. You have two different approaches:

The "One-Step Sprint" (Drift Models): You look at a blank canvas and, in one giant leap, try to paint the final portrait instantly. It’s incredibly fast, but because you didn’t take the time to sketch or refine details, the result might be a bit rough or blurry. You can’t really "fix" it later because you already finished the painting.
The "Slow Sculptor" (Flow Matching/Diffusion Models): You start with a block of clay (or noise) and slowly, step-by-step, refine it. You smooth it out, add details, and adjust the shape over many small steps. This takes a long time and requires a lot of computation, but the final result is usually very high quality. If you want it to look even better, you can just take more steps.

The Dilemma: Usually, you have to choose. Do you want it fast and okay? Or slow and perfect?

The Solution: Drift Flow Matching (DFM)

The authors of this paper created a new method called Drift Flow Matching (DFM). Think of DFM as a "Smart Shortcut" that gives you the best of both worlds.

It’s like having a GPS that can either:

Teleport you to your destination instantly (like the One-Step Sprint).
Guide you step-by-step along the best route (like the Slow Sculptor).

With DFM, you can choose how much "effort" you want to put in at the moment you need the image. If you’re in a hurry, you use one step. If you want the highest quality, you use many steps. The model is flexible enough to handle both.

How It Works: The "Group Chat" Analogy

To understand how DFM learns to do this, imagine a classroom where students are learning to move from Point A (Noise) to Point B (Perfect Image).

Old Way (Flow Matching):
The teacher tells each student exactly where to step next based on their specific position. "You are here, so move 2 inches to the right." It’s very precise but requires calculating every single tiny movement.

Old Way (Drift Models):
The teacher just says, "Everyone, move toward the target!" The students pull themselves toward the goal in one big jump. It’s fast, but they don’t learn the path, just the destination.

The DFM Way:
The teacher groups students based on where they are in their journey.

Group 1 is at the start.
Group 2 is halfway there.
Group 3 is near the end.

For each group, the teacher doesn’t just say "go to the end." Instead, the teacher says, "Look at where your group should be in 5 minutes. Move toward that spot."

If a student is at the start, they learn to jump toward the "halfway" spot.
If a student is halfway, they learn to jump toward the "near-end" spot.

By training the model to understand these short hops between any two points in time (not just start-to-finish), the model learns the entire path. This allows it to either take one giant hop (if you ask it to go from start to finish instantly) or many small hops (if you ask it to refine the image step-by-step).

Why Is This Useful?

The paper demonstrates that DFM is versatile. It works well for:

Making Images: Whether it’s simple digits (MNIST), realistic faces (FFHQ), or complex photos (ImageNet), DFM can generate them quickly or with high detail, depending on what you need.
Robotics: Imagine a robot arm trying to pick up a cup. DFM helps the robot decide its movements. It can act quickly for simple tasks or take more careful, refined steps for complex tasks like hanging a tool on a hook.

The Bottom Line

Drift Flow Matching bridges the gap between "fast but rough" and "slow but perfect." It teaches AI models to understand the journey from noise to image, not just the destination. This means you can dial up the quality by adding more steps, or dial it down for speed, all within the same model. It’s a flexible tool that adapts to how much time and computing power you have available.

Technical Summary: Drift Flow Matching

Problem Statement
Iterative generative models, such as Flow Matching and Diffusion models, exhibit strong test-time scaling behavior, where increasing the number of inference steps improves generation quality, controllability, and distribution coverage. In contrast, Drift Models enable efficient one-step generation by learning a direct transport map from a source to a target distribution. However, this direct generation paradigm lacks a natural mechanism for test-time scaling, limiting flexibility in balancing computational cost against generation quality. The paper identifies a gap between the efficiency of one-step Drift Models and the iterative refinement capabilities of multi-step Flow Matching methods.

Methodology
The authors propose Drift Flow Matching (DFM), a framework that integrates drifting generative modeling with flow-based iterative generation. DFM learns a two-time transport model capable of moving samples from any current time $t$ to any future time $r > t$ along a Flow Matching probability path from $p_0$ to $p_1$ .

Two-Time Marginal Construction: DFM utilizes the conditional path construction of Flow Matching, where a coupling $\pi(x_0, x_1)$ connects source $p_0$ and target $p_1$ . By sampling time pairs $(t, r)$ from a distribution $\rho$ on the domain $\Delta = \{(t, r) \in [0, 1]^2 : 0 \le t \le r \le 1\}$ , the method constructs marginal distribution pairs $(p_t, p_r)$ .
Mean-Velocity Parameterization: Instead of regressing pointwise conditional velocities as in standard Flow Matching, DFM parameterizes a mean velocity field $u_\theta(x_t, t, r)$ that transports the current distribution $p_t$ to a predicted distribution $q^\theta_{t,r}$ at time $r$ . The transport map is defined as $T^\theta_{t,r}(x_t) = x_t + (r-t)u_\theta(x_t, t, r)$ .
Drift-Based Supervision: The core innovation lies in the training objective. DFM applies drift supervision to push the predicted marginal $q^\theta_{t,r}$ toward the true target marginal $p_r$ . The drift field $V_{q^\theta_{t,r}, p_r}$ is computed using kernel-weighted averages, attracting samples toward the target distribution while subtracting self-correction terms from the current model distribution.
Group-wise Training: Crucially, because each time pair $(t, r)$ induces distinct marginal distributions, the drift field and loss must be evaluated group-wise within each time-pair bin during mini-batch training, rather than mixing samples across different time pairs.

Key Contributions

Unified Framework: DFM bridges the gap between one-step Drift Models and multi-step Flow Matching, allowing the model to behave like a Drift Model for one-step inference and like a Flow Matching model for multi-step inference.
Test-Time Scaling for Drift Models: The framework equips drift-based methods with test-time scaling capability, enabling generation quality to be refined through multiple inference steps when desired.
Marginal Transport Learning: Unlike Flow Matching’s pointwise regression, DFM directly learns transport between marginal distributions $(p_t, p_r)$ , offering a novel optimization perspective that preserves the efficiency of direct transport maps.

Results
Extensive experiments demonstrate the effectiveness and generality of DFM across various tasks:

Synthetic Data: Visualizations on 2D datasets (e.g., "F", "M", moons, checkerboard) show that DFM produces promising results with one-step generation and further improves sample quality and distribution coverage with increased inference steps.
Conditional Generation (MNIST & FFHQ): On MNIST, DFM achieves competitive EMD and accuracy compared to Flow Matching and Mean Flow methods. On FFHQ, DFM demonstrates flexible trade-offs between sampling speed and quality, with latent space visualizations showing improved distribution coverage as inference steps increase.
ImageNet-1k: In large-scale synthesis, DFM achieves competitive FID and IS scores with one-step generation (NFE=1) and shows clear test-time scaling, with quality improving as NFE increases to 2, 5, and 10.
Robotic Control: Applied to robotic control tasks (e.g., Lift, Can, ToolHang), DFM preserves the one-step inference capability of Drift Policy while exhibiting clear test-time scaling behavior, improving success rates with additional inference steps.

Significance
The paper claims that DFM establishes a more flexible generative paradigm where direct generation and iterative refinement are no longer separate design choices but can be viewed within a single framework. By unlocking the test-time scaling capability of Drift Models, DFM allows generation to adapt to different quality–efficiency requirements, providing a novel approach that combines the efficiency of direct transport maps with the scalability of iterative flow-based methods.

The Core Problem: Speed vs. Quality

The Solution: Drift Flow Matching (DFM)

How It Works: The "Group Chat" Analogy

Why Is This Useful?

The Bottom Line

More like this