Flow Matching

Implementation of Conditional Flow Matching for generative modeling.

Visualizations

Particle Flow Animation

Vector Field Evolution

Installation

pip install flow-matching

Or install from source:

git clone https://github.com/Niclas-Mettenleiter/flow-matching.git
cd flow-matching
pip install -e .

Usage

2D Example (Two Moons)

import torch
from sklearn.datasets import make_moons
from flow_matching import (
    ResNetVectorField,
    get_device,
    sample_2d_batch,
    train_flow_model,
    visualize_2d_flow,
)

# Setup
device = get_device()
data, _ = make_moons(n_samples=10000, noise=0.05)
train_tensor = torch.tensor(data, dtype=torch.float32).to(device)

# Define sampling function
def sample_batch(batch_size):
    return sample_2d_batch(train_tensor, batch_size, device, use_ot=True)

# Train model
model = ResNetVectorField(input_dim=2, n_hidden=256, time_embed_dim=32).to(device)
model, ema_model, losses = train_flow_model(sample_batch, model, device, total_steps=1000)

# Visualize
visualize_2d_flow(ema_model, device)

1D Example (Laplace Distribution)

import torch
from flow_matching import (
    VectorFieldMLP,
    Standardization,
    sample_1d_batch_with_transform,
    train_flow_model,
)

# Setup transform
transform = Standardization(mean=50.0, std=21.2).to(device)

# Define sampling
def sample_batch(batch_size):
    return sample_1d_batch_with_transform(
        source_dist=torch.randn,
        target_dist=torch.distributions.Laplace(50.0, 15.0),
        transform=transform,
        batch_size=batch_size,
        device=device,
    )

# Train
model = VectorFieldMLP(input_dim=1, n_hidden=64, time_embed_dim=16).to(device)
model, ema_model, losses = train_flow_model(sample_batch, model, device)

Examples

See the notebooks/ directory for detailed examples:

• flow_matching_two_moons.ipynb - 2D generative modeling
• flow_matching_laplace.ipynb - 1D distribution matching

Development

Install with development dependencies:

pip install -e ".[dev]"

Code Quality

# Format code
hatch run code-quality:format

# Run linting
hatch run code-quality:lint

# Run all checks
hatch run code-quality:check

License

flow-matching is distributed under the terms of the MIT license.