Interpreting the Geometric Structure of Flow Matching

Currently training a JiT-style ViT with x-prediction and v-loss on rotated class-conditional MNIST.

Setup

brew install uv
uv sync

Usage

# train (use --offline to skip wandb)
uv run python -m scripts.train --offline

# generate samples from a checkpoint
uv run python -m scripts.generate --experiment <name>

File Structure

geoflow/       # library (model, denoiser, checkpoint, optim, etc.)
scripts/       # entry points
  train.py
  generate.py