# 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.