| # EqFM on CIFAR-10, EDM2 U-Net, from scratch (DeltaAI, 4x GH200, global batch 512) |
|
|
| Code: github.com/justinlin26/EqFM, branch `mnist-smf-ablation` (`scripts/train_stage2_flowmap_mnist.py`, |
| `eqfm/models/edm2_flowmap.py`); recipe notes in `docs/RUN_ME_edm2_v2.md`. |
| Unconditional; the flow map is X_sigma(x) = x + sigma * v(x, sigma). Network: the flow-map paper's EDM2 CIFAR |
| U-Net (Boffi et al., arXiv 2505.18825), 54.67M params. Shifted MeanFlow loss with a live self-teacher. |
| |
| | Folder | Run | Recipe | |
| |---|---|---| |
| | `run_b_trunc_std1/` | `edm2_large_b512_trunc_std1_s0` | truncated clock a=0.9, noise std 1 (sigma_data 1); diag weight 768 (256 from step 211277); LR 4e-3, 8e-4 at 130k, 1e-4 at 250k | |
| | `run_a_linear/` | `edm2_large_b512_s0` | linear clock, noise std 0.5 (sigma_data 0.5); diag weight 768 (3072 from step 60298); LR 4e-3, cosine to 1e-4 from 242k | |
| |
| Files per run: |
| - `ema_step_XXXXXXX.pt`: best EMA checkpoint (EMA weights + training args). Run B: 125k. Run A: 300k. |
| - `latest_step_0402000.pt`: final full checkpoint (raw weights `model`, `ema`, AdamW `optimizer`, `step`) -- resumable. |
| - `args.json`: trainer arguments at the last launch; `log_fid.jsonl`: in-run FID-2k (EMA and raw) every 25k steps. |
|
|
| FID (CIFAR-10 train statistics, clean-fid legacy_tensorflow, expuniform sampler, EMA): |
| |
| | Checkpoint | Metric | 1 | 2 | 4 | 8 | 16 steps | |
| |---|---|---|---|---|---|---| |
| | Run A, 300k | FID-50k | 11.30 | 12.81 | 10.54 | 7.50 | 7.04 | |
| | Run A, 402k | FID-50k | 14.24 | 15.77 | 13.27 | 8.41 | 7.10 | |
| | Run B, 100k | FID-50k | 8.01 | 8.19 | 9.67 | 7.72 | 4.72 | |
| | **Run B, 125k** | **FID-50k** | **7.29** | **5.83** | **7.93** | **7.18** | **4.89** | |
| | Run B, 150k | FID-50k | 7.44 | 6.69 | 8.30 | 7.15 | 5.01 | |
| | Run B, 402k | FID-2k (in-run) | 25.0 | 20.7 | 22.1 | 20.9 | – | |
| |
| (In-run FID-2k reads about 13 points higher than FID-50k. FID-50k JSONs are in `fid50k/`.) |
| |
| Evaluate: `python scripts/eval_fid_cifar.py --mode eqfm --ckpt <file> --nfe 1 2 4 8 16 --schedule expuniform |
| --num-gen 50000 --batch-size 500 --seed 2024` (reads clock and noise std from the checkpoint). |
| Resume: full checkpoint via `--resume-from latest_step_0402000.pt`; from an EMA snapshot via |
| `--resume-from ema_step_XXXXXXX.pt --rewarmup 1000` (optimizer reset), into a new run directory. |
| |
