miniGen: A 2D Playground for Generative Models

Generative modeling is probability distribution mapping in essense.

This is a minimalist, modular playground for training and visualizing generative models on 2D synthetic data (source distribution → target distribution). Built for rapid experimentation, prototyping, and learning, it helps me stay current with new generative modeling ideas. The focus is on core algorithmic logic, which largely carries over to production systems; most real-world differences stem from scale and training refinements.

Visualization

Distribution Matching Process

Below are visualizations showing how distribution matching evolves during training across different generative modeling paradigms/methods.

Variational Autoencoder	GAN	Energy-Based Model	Autoregressive Model	Score Matching (Hyvärinen)
Score Matching (NCSN)	Diffusion (DDPM)	Normalizing Flow	Continuous NF	(Rectified) Flow Matching
Consistency Training	Shortcut Model	Improved MeanFlow	Alpha Flow	Drifting Model

The Score Fields (Score-based Models)

Visualizing the learned score fields over different training iterations.

Score Matching (Hyvärinen)

Score Matching (NCSN)

Diffusion (DDPM)

The Velocity Fields (Flow Models)

Visualizing the learned velocity fields over different training iterations.

Continuous Normalizing Flow

(Rectified) Flow Matching

Improved MeanFlow

The Drift Fields (Drifting Model)

Visualizing the "drifting" motion at inference time.

Swiss Roll

Gaussian Mixture

Checkerboard

The drifting model performs one-step generation:

• Blue dots represent randomly initialized points
• Red dots represent the drifted outputs of the model
• Green lines visualize their one-step drifting process or trajectories.

Usage

Training

Train any model by pointing to its config:

python train.py --config configs/flow_matching/FlowMatching_MLP_SwissRoll.yaml

Override config values from the command line:

python train.py --config configs/diffusion/DDPMSampler_MLP_SwissRoll.yaml training.total_steps=5000

Add --skip_save_ckpt to skip saving checkpoints.

Configs are organized as configs/{paradigm}/{Model}_{Arch}_{Dataset}.yaml. Each config specifies the model type, architecture, dataset, training hyperparameters, and inference settings.

Supported Models

Currently 18 generative modeling paradigms are supported — see paradgims/ for the full list.

Datasets

Nine 2D synthetic distributions: swiss_roll, gaussian_mixture, moons, circles, checkerboard, spirals, pinwheel, rings, s_curve.

Outputs

Training outputs are saved to logs/{timestamp}_{exp_name}/:

• train.log — training metrics (step, loss, lr, throughput)
• saved_images/ — sample visualizations at regular intervals
• saved_checkpoints/ — model checkpoints
• visualize.html — interactive image gallery

Quick Start

This project uses uv for environment management.

uv sync

This installs all dependencies (PyTorch, NumPy, matplotlib, OmegaConf, tqdm) into a local .venv. The PyTorch wheel is platform-aware: CUDA-enabled on Linux (x86_64), CPU+MPS on macOS (Apple Silicon).

Then run training with:

uv run python train.py --config configs/flow_matching/FlowMatching_MLP_SwissRoll.yaml

Requirements

• Python 3.10+
• uv (recommended) or pip
• PyTorch, NumPy, matplotlib, OmegaConf, tqdm (managed by pyproject.toml)