tinyflow: Flow Matching using Tinygrad

An implementation of Flow Matching algorithms using the tinygrad library. This project prioritizes pedagogical clarity and extensibility, for learning, research prototyping, and understanding flow-based generative models.

What is Flow Matching?

Flow matching is an approach to generative modeling that learns to transform a noise distribution into a data distribution through a continuous-time flow. The model learns the velocity field of the transformation by direct regression, and samples are generated by solving an ODE from t=0 (noise) to t=1 (data).

Features

• Paths: AffinePath, OptimalTransportPath
• Schedulers: Linear, Cosine, Polynomial, Linear Variance Preserving (LVP)
• ODE solvers: Euler, Heun, Midpoint, RK4, DDIM (deterministic, same result as Euler)
• Time samplers for training: uniform, logit-normal, mode/heavy-tail, cos
• Models: MLPNetwork, UNetMNIST, UNetCIFAR10, UNetCIFAR10Large (with self-attention)
• Datasets: MNIST, Fashion MNIST, CIFAR-10, 2D moons
• LR schedulers: step, cosine, warmup
• Experiment tracking: Hydra configuration and MLflow logging
• Generation tools: static sample grids and animated GIFs

Quick Start

Installation

Requires Python 3.12 or later and uv.

git clone https://github.com/iliailmer/tinyflow
cd tinyflow
uv sync

Datasets

# Fashion MNIST and CIFAR-10
uv run scripts/download_datasets.py --all

• Moons: generated in code, no download.
• MNIST: not downloaded by the script. Place JPG files in dataset/mnist/trainset/trainingSet/<digit>/*.jpg.
• Fashion MNIST: idx files in dataset/fashion_mnist/.
• CIFAR-10: batch files in dataset/cifar10/cifar-10-batches-py/.

Training

# 2D moons (fast, for testing)
uv run scripts/train.py dataset=moons model=neural_network training=moons_default scheduler=linear

# MNIST (default config)
uv run scripts/train.py

# Fashion MNIST
uv run scripts/train.py dataset=fashion_mnist

# CIFAR-10
uv run scripts/train.py dataset=cifar10 model=unet training=cifar10_default

Image training runs for 1 epoch by default. Set training.num_epochs=N for more. One epoch checks that the pipeline runs; usable CIFAR-10 samples need many more.

Weights are saved to weights/model_<dataset>_<model>_ts_<time_sampler>_<scheduler>.safetensors, for example weights/model_mnist_unet_ts_uniform_linear.safetensors.

Generating Samples

Pass the weights path as generation.model_path. The scheduler is used only in training; generation needs the weights and a solver.

# Static grid for MNIST
uv run scripts/generate.py generation.model_path=weights/model_mnist_unet_ts_uniform_linear.safetensors

# Animated GIF of the flow
uv run scripts/generate.py generation.model_path=weights/model_mnist_unet_ts_uniform_linear.safetensors --animated

# CIFAR-10
uv run scripts/generate.py dataset=cifar10 generation.model_path=weights/model_cifar10_unet_ts_uniform_linear.safetensors

# 2D moons (switch the base config)
uv run scripts/generate.py --config-name=generate_moons_config generation.model_path=weights/model_moons_neural_network_ts_uniform_linear.safetensors

# Grid size and number of ODE steps
uv run scripts/generate.py generation.model_path=<path> generation.grid_size=3 generation.num_steps=50

Outputs are saved in outputs/generated/.

Configuration with Hydra

# Override parameters
uv run scripts/train.py scheduler=cosine optimizer.lr=0.01

# Compare schedulers (multirun)
uv run scripts/train.py -m scheduler=linear,cosine,polynomial

# Use an experiment config
uv run scripts/train.py +experiment=quick_test

Group	Options	Default
dataset	mnist, fashion_mnist, cifar10, moons	mnist
model	unet, unet_large (CIFAR-10 only), neural_network	unet
scheduler	linear, cosine, polynomial, lvp	linear
solver	euler, heun, midpoint, rk4, ddim	heun
time_sampler	uniform, logit_normal, mode_heavy_tail, cos	uniform
lr_scheduler	none, step, cosine, cosine_warm_restarts, cosine_warmup, cosine_warmup_long, step_warmup	none
training	mnist_default, cifar10_default, cifar10_improved, moons_default	mnist_default
experiment	quick_test, moons_comparison, mnist_cosine, mnist_polynomial	none

Sampling API

from tinygrad.tensor import Tensor
from tinyflow.solver import Heun, uniform_time_grid
from tinyflow.utils import preprocess_time

solver = Heun(model, preprocess_hook=preprocess_time)
x = solver.solve(Tensor.randn(16, 1, 28, 28), uniform_time_grid(20))

solve(x, time_grid) integrates over a grid of N + 1 time points. Pass any increasing grid for a non-uniform schedule.

Experiment Tracking with MLflow

Each training run is logged to mlflow.db (SQLite) in the repository root.

uv run mlflow ui --backend-store-uri sqlite:///mlflow.db
# View at http://localhost:5000

MLflow records:

• Parameters: the full config, with dotted keys such as optimizer.lr
• Metrics: loss, learning rate
• Artifacts: loss curve, generated samples
• Tags: dataset, model, scheduler, solver, time sampler

Project Structure

tinyflow/
├── nn.py                      # Models (MLP, UNets)
├── losses.py                  # Loss functions
├── trainer.py                 # Training loop with MLflow logging
├── dataloader.py              # Dataset loaders (MNIST, Fashion MNIST, CIFAR-10)
├── time_sampler.py            # Training time samplers
├── utils.py                   # Visualization and time preprocessing
├── path/
│   ├── path.py                # Flow matching paths (Affine, OT)
│   └── scheduler.py           # Schedulers (Linear, Cosine, Polynomial, LVP)
├── solver/
│   ├── solver.py              # Base ODE solver, solve(), uniform_time_grid
│   ├── euler.py               # Euler method
│   ├── heun.py                # Heun's method
│   ├── midpoint.py            # Midpoint method
│   ├── rk4.py                 # 4th order Runge-Kutta
│   └── ddim.py                # DDIM-style deterministic sampling
└── nn_utils/
    ├── conv.py                # Convolutional building blocks
    ├── attention.py           # Self-attention
    ├── time_embedding.py      # Time embedding layers
    └── lr_scheduler.py        # Learning rate schedulers

configs/                       # Hydra configuration files
├── config.yaml                # Training config
├── generate_config.yaml       # Image generation config
├── generate_moons_config.yaml # Moons generation config
├── cifar10_improved.yaml      # CIFAR-10 config with the large UNet
└── dataset/ model/ scheduler/ solver/ time_sampler/ lr_scheduler/
    optimizer/ training/ generation/ experiment/

scripts/
├── _common.py                 # Shared Hydra config helpers
├── train.py                   # Training entry point (moons and image datasets)
├── generate.py                # Generation entry point (static grid / animation)
└── download_datasets.py       # Dataset download helper

experiments/                   # ODE-solver and time-schedule research scripts
├── solver_compare.py          # Compare solvers (Euler/Heun/RK4) at fixed NFE
├── schedule_compare.py        # Compare uniform vs back-loaded time schedules
├── curvature_profile.py       # Velocity-field curvature / LTE-proxy profiling
├── adaptive_sampler.py        # Curvature-adaptive step size
└── homotopy_runner.py         # Dispatcher for the homotopy experiments

tests/                         # pytest tests

Development

# Tests
uv run pytest tests/ -v

# Lint, format check, type check (same as CI)
uv run ruff check tinyflow/ tests/ scripts/ experiments/
uv run ruff format --check tinyflow/ tests/ scripts/ experiments/
uv run pyright tinyflow/

# Pre-commit hooks
uv run pre-commit install

Planned Work

• Exponential Moving Average (EMA) of the weights
• Stronger time conditioning in the UNets
• Class-conditional generation
• Pre-trained models

Contributing

Contributions welcome. This project prioritizes:

• Educational clarity over performance
• Clean, readable code
• Extensibility for research

Why Tinygrad?

• Lightweight: few dependencies
• Educational: a smaller codebase than PyTorch or JAX
• Research-friendly: easy to modify for experiments

References

Core Papers

• Flow Matching for Generative Modeling - Lipman et al., 2022
• Flow Matching Guide and Code - Lipman et al., 2024

Frameworks

• Tinygrad - ML framework
• Hydra - Configuration management
• MLflow - Experiment tracking

License

MIT License - See LICENSE for details.