| # TODO.md — Next Steps for NSGF++ Reproduction |
|
|
| ## Current Status |
|
|
| | Experiment | Pool Building | Phase 1 (NSGF) | Phase 2 (NSF) | Phase 3 (Predictor) | Inference | Eval | |
| |-----------|:---:|:---:|:---:|:---:|:---:|:---:| |
| | **2D 8gaussians** | ✅ | ✅ | — | — | ✅ | ✅ W2=2.04 (small run) | |
| | **MNIST** | ✅ | 🔶 runs, loss converging (~0.03), interrupted at 9.5K/100K | untested on GPU | untested on GPU | untested | untested | |
| | **CIFAR-10** | 🔶 OOM fixed (batch 128→32), untested on GPU | untested | untested | untested | untested | untested | |
|
|
| ✅ = verified working 🔶 = partially done ❌ = blocked |
|
|
| --- |
|
|
| ## Immediate — Run Full Experiments |
|
|
| ### 1. MNIST full run on T4 |
|
|
| The most important next step. All code bugs are fixed. Need a clean Kaggle run. |
|
|
| ```bash |
| cd /kaggle/working/ && rm -rf nsgf-plusplus |
| git clone https://huggingface.co/rogermt/nsgf-plusplus |
| cd nsgf-plusplus && pip install -r requirements.txt |
| |
| # Phase 1: pool (~7 min) + NSGF training (100K steps, ~2.5 hrs) |
| python main.py --experiment mnist |
| |
| # If session runs out, next session: |
| python main.py --experiment mnist --resume-phase 2 |
| |
| # If Phase 2 done: |
| python main.py --experiment mnist --resume-phase 3 |
| ``` |
|
|
| **Expected runtimes on T4:** |
| - Pool building (1500 batches): ~7 min |
| - Phase 1 NSGF (100K steps): ~2.5 hours |
| - Phase 2 NSF (100K steps): ~3-4 hours (each step does NSGF inference + NSF forward/backward) |
| - Phase 3 Predictor (40K steps): ~1.5 hours |
| - **Total: ~7-8 hours** — tight for one 9-hour Kaggle session |
|
|
| **Alternative: use `--train-iters 50000` for Phase 1+2 to fit in one session, accept lower quality.** |
|
|
| **Paper target: FID ≈ 3.8 at NFE=60** |
|
|
| --- |
|
|
| ### 2. CIFAR-10 first test on T4 |
|
|
| After MNIST works, test CIFAR with reduced Sinkhorn batch. |
|
|
| ```bash |
| # Smoke test first (should run ~2 min) |
| python main.py --experiment cifar10 --pool-batches 10 --train-iters 50 |
| |
| # If smoke test passes, real Phase 1: |
| python main.py --experiment cifar10 --train-iters 50000 |
| |
| # Subsequent sessions: |
| python main.py --experiment cifar10 --resume-phase 2 --train-iters 50000 |
| python main.py --experiment cifar10 --resume-phase 3 |
| ``` |
|
|
| **If still OOMs**: try `--sinkhorn-batch 16 --pool-batches 20000` |
|
|
| **Paper target: FID ≈ 5.55, IS ≈ 8.86 at NFE=59** |
|
|
| --- |
|
|
| ### 3. 2D full-scale run |
|
|
| Quick win to validate against paper numbers. Should take ~20 min on T4. |
|
|
| ```bash |
| python main.py --experiment 2d --dataset 8gaussians --steps 10 |
| ``` |
|
|
| **Paper target: W2 ≈ 0.285 for 8gaussians** |
|
|
| Current small-run W2=2.04 is expected — only used 10 pool batches + 1000 iters. Full run (200 batches, 20K iters) should drop dramatically. |
|
|
| Also run other 2D datasets: |
| ```bash |
| python main.py --experiment 2d --dataset moons --steps 10 |
| python main.py --experiment 2d --dataset scurve --steps 10 |
| python main.py --experiment 2d --dataset checkerboard --steps 10 |
| ``` |
|
|
| --- |
|
|
| ## Medium-term — Code Improvements |
|
|
| ### 4. Step-level resume within phases |
|
|
| Current `--resume-phase` skips completed phases but restarts the current phase from step 0. For 100K-step phases, mid-phase interruption still loses progress. Need: |
| - Load `nsgf_checkpoint.pt` / `nsf_checkpoint.pt` / `predictor_checkpoint.pt` |
| - Resume optimizer state + step counter |
| - Continue from last checkpoint step |
|
|
| ### 5. EMA (Exponential Moving Average) for image models |
|
|
| Paper uses EMA for MNIST and CIFAR-10 (standard in diffusion/flow models). Current code doesn't implement EMA. This likely affects FID significantly. |
|
|
| ### 6. Learning rate scheduler |
|
|
| Paper may use cosine decay or warmup. Currently using constant lr. Check if this matters for convergence. |
|
|
| ### 7. FID evaluation correctness |
|
|
| Verify that `evaluation.py`'s FID computation matches the standard protocol: |
| - InceptionV3 features from `pool3` layer (2048-dim) |
| - 10K generated vs 10K test samples |
| - Proper image preprocessing (resize to 299×299 for Inception) |
| - Compare against `pytorch-fid` or `clean-fid` for sanity check |
|
|
| ### 8. Inception Score evaluation |
|
|
| Implement properly for CIFAR-10 if not already correct. Paper reports IS=8.86. |
|
|
| --- |
|
|
| ## Longer-term — Towards Paper Numbers |
|
|
| ### 9. Full paper hyperparameters |
|
|
| Once code is stable, run with exact paper configs (no iteration reduction): |
| - MNIST: 100K + 100K + 40K iterations |
| - CIFAR-10: 200K + 200K + 40K iterations |
| - This requires A100 or multiple Kaggle sessions with checkpointing |
|
|
| ### 10. Ablation: NSGF vs NSGF++ |
|
|
| Run NSGF-only (Phase 1 only, no straight flow) and compare FID/W2 against NSGF++ to verify the two-phase approach actually helps. Paper shows clear improvement. |
|
|
| ### 11. NFE sweep |
|
|
| Paper reports results at various NFE (number of function evaluations). Test: |
| - MNIST: NFE = 10, 20, 40, 60 |
| - CIFAR: NFE = 10, 20, 40, 59 |
| - Compare FID vs NFE curve against paper's Figure 3 |
|
|
| ### 12. pykeops for faster Sinkhorn |
|
|
| Install `pykeops` to enable geomloss `online` backend. This avoids materializing the full N×N cost matrix and should be much faster + lower VRAM for image experiments. Could enable using paper's original batch_size=128 on T4. |
| |
| ```bash |
| pip install pykeops |
| # Then in config or code: |
| # backend: "online" instead of "tensorized" |
| ``` |
| |
| --- |
| |
| ## Known Limitations |
| |
| - **Single-GPU only** — no DDP, T4×2 wastes one GPU |
| - **No EMA** — standard in flow/diffusion, likely hurts FID |
| - **No mixed precision** — fp32 only, could halve VRAM with fp16/bf16 |
| - **No gradient accumulation** — batch size is hard-limited by VRAM |
| - **Kaggle checkpoint persistence** — checkpoints lost between sessions unless manually saved |
| |
