# Flow Matching

## Core Idea

Flow Matching learns a time dependent velocity field $v_\theta(\mathbf{x}, t)$ that transports samples from a simple prior distribution to a target data distribution.

### Distributions and Path

- Prior (noise):
  - $\mathbf{x}_0 \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$
- Target (data):
  - $\mathbf{x}\_1 \sim p\_\text{data}$ (two moons samples)

The code uses a linear interpolation path between noise and data with $t \sim \mathcal{U}(0, 1)$:

```math
\mathbf{x}_t = (1 - t)\mathbf{x}_0 + t\mathbf{x}_1.
```

Taking the derivative with respect to time:

```math
\frac{\mathrm{d}\mathbf{x}_t}{\mathrm{d}t} = \mathbf{x}_1 - \mathbf{x}_0.
```

So the training target velocity used in this project is:

```math
\mathbf{v}_\text{target} = \mathbf{x}_1 - \mathbf{x}_0.
```

### Training Objective

The model is trained with mean squared error (MSE):

```math
\mathcal{L}(\theta) = \mathbb{E}\left[\left\lVert v_\theta(\mathbf{x}_t, t) - (\mathbf{x}_1 - \mathbf{x}_0)\right\rVert_2^2\right].
```

In code:

- Sample a mini-batch of $\mathbf{x}_1$ from the two moons dataset (`data_samples`)
- Sample a matching mini-batch of $\mathbf{x}_0 \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ (`noise_samples`)
- Sample $t \sim \mathcal{U}(0, 1)$ (`time_values`)
- Compute $\mathbf{x}_t = (1-t)\mathbf{x}_0 + t\mathbf{x}_1$ (`interpolated_points`)
- Compute $\mathbf{v}_\text{target} = \mathbf{x}_1 - \mathbf{x}_0$ (`true_velocity`)
- Predict $\hat{\mathbf{v}} = v\_\theta(\mathbf{x}\_t, t)$ (`predicted_velocity`) and minimize MSE against $\mathbf{v}\_{\mathrm{target}}$ (`true_velocity`)

### Sampling (Generation)

After training, samples are generated by integrating the learned velocity field from $t = 0$ to $t = 1$.

This project uses forward Euler integration with $N$ time steps:

```math
\Delta t = \frac{1}{N},\quad \mathbf{x}_0 \sim \mathcal{N}(\mathbf{0}, \mathbf{I}),\quad t_k = \frac{k}{N}.
```

For $k = 0, 1, \dots, N-1$:

```math
\mathbf{x}_{t_{k+1}} = \mathbf{x}_{t_k} + v_\theta(\mathbf{x}_{t_k}, t_k)\,\Delta t.
```
