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.