A comprehensive implementation of Flow Matching algorithms using the tinygrad library. This project prioritizes pedagogical clarity and extensibility, making it ideal for learning, research prototyping, and understanding flow-based generative models.
Flow matching is a modern approach to generative modeling that learns to transform noise distributions into data distributions through continuous-time flows. Unlike diffusion models, flow matching directly learns the velocity field of the transformation, often leading to:
- Faster sampling (fewer integration steps needed)
- Simpler training (direct regression instead of score matching)
- Better mathematical clarity (continuous normalizing flows)
- Multiple Schedulers: Linear, Cosine, Polynomial, Linear Variance Preserving (LVP)
- ODE Solvers: Euler, Midpoint, RK4, Heun, DDIM
- Architectures: MLP networks, UNet for images
- Datasets: MNIST, Fashion MNIST, CIFAR-10, 2D toy datasets (moons)
- Time Embeddings: Sinusoidal/Fourier embeddings for temporal conditioning
- Metrics: FID (Fréchet Inception Distance) with trained classifiers
- Experiment Tracking: Hydra configuration + MLflow logging
- Generation Tools: Class predictions, FID computation, animated generation
- FID Metric: Measure generation quality quantitatively
- Classifier Predictions: Show predicted class and confidence on generated images
- Animated Generation: Visualize the flow from noise to data
- MLflow Integration: Track experiments, hyperparameters, and metrics
# Clone repository
git clone https://github.com/iliailmer/flow_matching_tinygrad
cd flow_matching_tinygrad
# Install dependencies (recommended: use uv)
uv sync# Train on 2D moons dataset (fast, for testing)
uv run scripts/train.py dataset=moons model=neural_network training=moons_default scheduler=linear
# Train on MNIST
uv run scripts/train.py
# Train on Fashion MNIST
uv run scripts/train.py dataset=fashion_mnist
# Train on CIFAR-10
uv run scripts/train.py dataset=cifar10 model=unet# Generate static grid for MNIST
uv run scripts/generate.py generation.model_path=model_mnist_unet_linear.safetensors
# Generate with class predictions and FID score
uv run scripts/generate.py \
generation.model_path=model_mnist_unet_linear.safetensors \
generation.show_predictions=true \
generation.compute_fid=true
# Generate animated GIF showing the flow process
uv run scripts/generate.py \
generation.model_path=model_mnist_unet_linear.safetensors \
--animated
# Generate for 2D moons dataset (switch the base config)
uv run scripts/generate.py --config-name=generate_moons_config generation.model_path=model_moons_neural_network_linear.safetensors# Train classifier for MNIST
uv run scripts/train_fid_classifier.py --dataset mnist --epochs 10
# Train classifier for Fashion MNIST
uv run scripts/train_fid_classifier.py --dataset fashion_mnist --epochs 10
# Train classifier for CIFAR-10
uv run scripts/train_fid_classifier.py --dataset cifar10 --epochs 10Generated outputs are saved in outputs/generated/ directory.
tinyflow/
├── nn.py # Neural network architectures (UNet, MLP)
├── losses.py # Loss functions
├── trainer.py # Training loops with MLflow integration
├── dataloader.py # Dataset loaders (MNIST, Fashion MNIST, CIFAR-10)
├── metrics.py # FID metric and classifiers
├── utils.py # Visualization utilities
├── path/
│ ├── path.py # Flow matching paths (Affine, OT)
│ └── scheduler.py # Schedulers (Linear, Cosine, Polynomial, LVP)
├── solver/
│ ├── solver.py # Base ODE solver
│ ├── euler.py # Euler method
│ ├── rk4.py # 4th order Runge-Kutta
│ ├── midpoint.py # Midpoint method
│ ├── heun.py # Heun's method
│ └── ddim.py # DDIM-style deterministic sampling
└── nn_utils/
├── conv.py # Convolutional building blocks
└── time_embedding.py # Time embedding layers
configs/ # Hydra configuration files
├── config.yaml # Main configuration
├── model/ # Model architectures
├── scheduler/ # Scheduler types
├── optimizer/ # Optimizer settings
├── dataset/ # Dataset configurations
├── training/ # Training parameters
└── generation/ # Generation settings
scripts/
├── train.py # Unified training entry point (moons + image datasets)
├── generate.py # Unified generation entry point (predictions/FID/animation)
├── train_fid_classifier.py # Train FID classifier with validation
└── 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
docs/
├── MLFLOW_HYDRA_GUIDE.md # Experiment tracking guide
├── TRAINING_OPTIMIZATION.md # Performance optimization
├── ODE_SOLVER_IMPROVEMENTS.md # Solver improvements
└── CIFAR10_IMPLEMENTATION.md # CIFAR-10 setup guide
This project uses Hydra for flexible configuration management:
# Override specific parameters
uv run scripts/train.py scheduler=cosine optimizer.lr=0.001
# Compare multiple schedulers (multirun)
uv run scripts/train.py -m scheduler=linear,cosine,polynomial
# Use experiment configs
uv run scripts/train.py +experiment=quick_testSee docs/MLFLOW_HYDRA_GUIDE.md for detailed usage.
All experiments are automatically logged to MLflow:
# Start MLflow UI
uv run mlflow ui
# View at http://localhost:5000MLflow tracks:
- Hyperparameters (model, scheduler, optimizer settings)
- Metrics (loss, FID score)
- Artifacts (generated images, loss curves)
- Tags (model type, dataset)
This project is under active development. See ROADMAP.md for:
- Planned features (EMA, conditional generation, advanced samplers)
- Implementation priorities
- Research directions (discrete flows, manifold flows)
- Community building efforts
High Priority Next Steps:
- Exponential Moving Average (EMA) for better quality
- Class-conditional generation
- Comprehensive documentation and tutorials
- Pre-trained models
- ROADMAP.md - Development roadmap and priorities
- CLAUDE.md - Instructions for Claude Code assistant
- docs/MLFLOW_HYDRA_GUIDE.md - Experiment tracking guide
- docs/TRAINING_OPTIMIZATION.md - Performance optimization tips
- docs/ODE_SOLVER_IMPROVEMENTS.md - Advanced solver implementations
- docs/CIFAR10_IMPLEMENTATION.md - CIFAR-10 setup guide
# Run all tests
uv run pytest tests/ -v
# Run specific test file
uv run pytest tests/test_paths.py -v
# Run with coverage
uv run pytest tests/ --cov=tinyflow --cov-report=htmlContributions welcome! This project prioritizes:
- Educational clarity over performance
- Clean, readable code
- Comprehensive documentation
- Extensibility for research
See ROADMAP.md for planned features and priorities.
- Lightweight: Minimal dependencies, easier to understand
- Educational: Simpler codebase than PyTorch/JAX
- Transparent: Mathematical operations more visible
- Research-friendly: Easy to modify for novel experiments
- Growing ecosystem: Help build the tinygrad ML community
- Flow Matching for Generative Modeling - Lipman et al., 2022
- Flow Matching Guide and Code - Pooladian et al., 2024
MIT License - See LICENSE for details.
This project builds on the flow matching literature and the tinygrad ecosystem. Special thanks to the authors of the flow matching papers and the tinygrad community.