Thread (21 posts)

I'll walk you through standard flow matching step by step, building from the ground up.

1. The Core Idea

Flow matching learns to transport samples from a simple noise distribution to a complex data distribution by learning a vector field that guides this transformation.

Think of it like learning the "flow" of a river that carries particles from one location to another.

2. Mathematical Setup

Goal: Transform noise $x_0 \sim \mathcal{N}(0,I)$ into data $x_1 \sim q_{data}$

Key objects:

• Probability path: $p_t(x)$ - how the distribution changes over time $t \in [0,1]$
• Vector field: $u_t(x)$ - the "flow direction" at each point and time
• Flow map: $\phi_t(x)$ - where a particle starting at $x$ ends up at time $t$

3. The Relationship Between These Objects

The vector field $u_t(x)$ defines how particles move: $\frac{d\phi_t(x)}{dt} = u_t(\phi_t(x))$

This flow map pushes the initial distribution $p_0$ to create the probability path: $p_t = [\phi_t]_* p_0$

4. The Flow Matching Objective

Problem: We want to learn a neural network $v_t(x; \theta)$ that approximates the true vector field $u_t(x)$.

Objective: $\mathcal{L}_{FM}(\theta) = \mathbb{E}_{t \sim U[0,1], x \sim p_t(x)} \|v_t(x; \theta) - u_t(x)\|^2$

Issue: We don't know $p_t(x)$ or $u_t(x)$ - they're what we're trying to learn!

5. The Conditional Flow Matching Solution

Key insight: Build the marginal path from simpler conditional paths.

For each data point $x_1$, define a conditional probability path: $p_t(x|x_1) = \mathcal{N}(x | \mu_t(x_1), \sigma_t(x_1)^2 I)$

With boundary conditions:

• $p_0(x|x_1) = \mathcal{N}(x|0, I)$ (standard noise)
• $p_1(x|x_1) = \mathcal{N}(x|x_1, \sigma_{min}^2 I)$ (concentrated around data)

6. Constructing the Marginal Path

The marginal path is the mixture: $p_t(x) = \int p_t(x|x_1) q(x_1) dx_1$

The marginal vector field is: $u_t(x) = \int u_t(x|x_1) \frac{p_t(x|x_1)q(x_1)}{p_t(x)} dx_1$

7. The Conditional Vector Field

For the Gaussian conditional path, the conditional vector field is: $u_t(x|x_1) = \frac{\sigma_t'(x_1)}{\sigma_t(x_1)}(x - \mu_t(x_1)) + \mu_t'(x_1)$

8. The Tractable Training Objective

Conditional Flow Matching (CFM): $\mathcal{L}_{CFM}(\theta) = \mathbb{E}_{t \sim U[0,1], x_1 \sim q(x_1), x \sim p_t(x|x_1)} \|v_t(x; \theta) - u_t(x|x_1)\|^2$

Key theorem: $\nabla_\theta \mathcal{L}_{CFM} = \nabla_\theta \mathcal{L}_{FM}$

So we can optimize the tractable CFM objective instead!

9. Concrete Example: Optimal Transport Path

Choose simple linear interpolation:

• $\mu_t(x_1) = t \cdot x_1$
• $\sigma_t(x_1) = 1 - (1-\sigma_{min})t$

This gives: $u_t(x|x_1) = \frac{x_1 - (1-\sigma_{min})x}{1-(1-\sigma_{min})t}$

The training objective becomes: $\mathcal{L}_{CFM}(\theta) = \mathbb{E}_{t, x_1 \sim q(x_1), x_0 \sim \mathcal{N}(0,I)} \|v_t(\psi_t(x_0); \theta) - (x_1 - (1-\sigma_{min})x_0)\|^2$

where $\psi_t(x_0) = (1-(1-\sigma_{min})t)x_0 + tx_1$.

10. Training Algorithm

for batch in dataloader:
    # Sample time uniformly
    t = torch.rand(batch_size)
    
    # Sample noise
    x_0 = torch.randn_like(batch)
    
    # Create noisy samples
    x_t = (1 - (1-sigma_min)*t) * x_0 + t * batch
    
    # Compute target vector field
    target = batch - (1-sigma_min) * x_0
    
    # Predict vector field
    pred = model(x_t, t)
    
    # Compute loss
    loss = mse_loss(pred, target)
    loss.backward()

11. Sampling/Inference

To generate samples:

1. Start with noise: $x_0 \sim \mathcal{N}(0,I)$
2. Solve ODE: $\frac{dx}{dt} = v_t(x; \theta)$ from $t=0$ to $t=1$
3. The result $x_1$ is your generated sample

12. Why This Works

The beauty is that:

• We never need to know the complex marginal distribution $p_t(x)$
• We only work with simple Gaussian conditionals $p_t(x|x_1)$
• The conditional paths automatically "average out" to create the right marginal behavior
• The learned vector field captures the optimal transport between noise and data

This gives you a principled way to learn complex generative models that can produce high-quality samples with deterministic sampling (no randomness needed at inference time).