sayakpaul/nanoDiT-JAX-butterflies

nanoDiT-JAX · Butterflies

Final weights from a 600-epoch, 30,000-update run of nanoDiT-JAX, a JAX/Flax port of nanoDiT. Architecture adapted from DiT. Trained from scratch for educational image generation.

129,586,956 parameters; 64×64 RGB; 2×2 patches; width 768; 12 blocks; 12 attention heads; 75 classes. Dataset: 6,499 butterfly images. Batch 128, AdamW at 1e-4, weight decay 0.01, no learning-rate warmup or gradient clipping. One v6e TPU, bf16 computation, float32 parameters and optimizer state.

EMA was reset to live weights at step 8,750 (epoch 175), then updated with decay capped at 0.999. The earlier 0.9999 EMA was too slow for this run. This changed only the evaluation weights.

Final epoch-mean training MSE: 0.138425. Final 10-epoch trailing mean: 0.139209. These are training losses; no held-out or FID evaluation was performed.

Inference

On a TPU VM, or use the Colab notebook:

git clone https://github.com/sayakpaul/nanoDiT-JAX.git
cd nanoDiT-JAX
uv venv --python 3.12
uv pip install --python .venv/bin/python -r requirements.txt
.venv/bin/python inference.py --repo-id sayakpaul/nanoDiT-JAX-butterflies \
  --class-ids 0 1 2 3 4 --samples-per-class 4 --seed 43 --output runs/generated.png

The script downloads model.msgpack and config.json, then runs 150 Euler steps with CFG 5. List class IDs with .venv/bin/python inference.py --list-classes. For CPU/GPU, install the appropriate JAX build instead of the TPU-specific requirements.

The exported EMA weights reproduced the final TPU sample grid exactly. Inference code: commit 4ecbad6.

Files

• model.msgpack: float32 EMA parameters for inference (~518 MB).
• config.json: architecture, training settings and ordered class names.
• checkpoint.msgpack: full resumable state, including live weights, EMA, AdamW state, RNG and step (~2.07 GB).
• checksums.json: SHA-256 checksums of the weights, full checkpoint and config.
• metrics.jsonl, events.jsonl, training_config.json, training_run.json, summary.json: training measurements and provenance.

Training progress

Live-weight samples at epochs 10–600, matching the original repository's grid layout and frame rate: