Flow Matching: Unified Framework for Generative Modeling

Module: Generative AI | Difficulty: Advanced

Flow Matching

Learn velocity field such that:

interpolates between noise and data.

Conditional Flow Matching

where .

Rectified Flow

1. Train flow on pairs
2. ODE solve to get pairs
3. Re-train on rectified pairs
4. Repeat

Connection to Diffusion

For VP SDE:

import torch, torch.nn as nn

class RectifiedFlow(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    def forward(self, x0, x1):
        t = torch.rand(x0.size(0), device=x0.device).view(-1,1,1,1)
        xt = (1-t)*x0 + t*x1
        target = x1 - x0
        return ((self.model(xt, t) - target)**2).mean()
    def sample(self, n, steps=50):
        x = torch.randn(n, 3, 64, 64)
        dt = 1.0/steps
        for i in range(steps):
            t = torch.full((n,1,1,1), i/steps, device=x.device)
            x = x + self.model(x, t)*dt
        return x

Research Insight: Rectified flows create straighter ODE trajectories, enabling fewer sampling steps. The key insight is that straight paths are easier to learn than curved ones.