Score-Based Diffusion Models for Undersampled MRI Reconstruction
Reimplemented an ICLR 2022 method from JAX in PyTorch to within 0.15 dB of the published results — then found the failure mode the original work did not characterise, and a 3-minute fine-tune that repairs it.
Reproduces Song et al., ICLR 2022 (arXiv:2111.08005) — a PyTorch reimplementation of a JAX-only method, with original uncertainty-quantification and failure-mode extensions on top. The pretrained NCSN++ prior is the original authors’; it was converted, not retrained.
Results
Every value links to the committed artifact that produces it.
| Metric | Value | Evidence |
|---|---|---|
| Reproduction of published PSNR at 4× / 8× / 24× (n = 256 slices) | ≤ 0.15 dB | REPRODUCTION.md |
| MMSE averaging gain over a single sample (N = 20) | +1.4 to +2.0 dB | REPRODUCTION.md |
| Uncertainty ↔ true error, rank correlation (Spearman)conservative bound; thesis reports up to 0.87 | 0.81–0.84 | thesis/content/ch05_results.tex |
| Fine-tune repair on 100% of 48 images (p = 7e-15) | +5.1 / +7.1 / +6.6 dB | REPRODUCTION.md |
| Flax → PyTorch weight conversion | 473 / 473 tensors | REPRODUCTION.md |
| Full reproduction cost (RTX 4090 spot) | ≈ 2.5 GPU-h · ≈ $1.70 | REPRODUCTION.md |
Reproduction of Table 1
Reproduction of Table 1 (Song et al., ICLR 2022) — n = 256 slices
| R | Ours (PSNR / SSIM) | Paper (PSNR / SSIM) | Zero-filled (PSNR / SSIM) |
|---|---|---|---|
| 4× | 40.03 / 0.967 | 39.91 / 0.965 | 28.64 / 0.759 |
| 8× | 37.75 / 0.960 | 37.63 / 0.958 | 26.28 / 0.748 |
| 24× | 29.57 / 0.881 | 29.42 / 0.880 | 20.79 / 0.515 |
source: REPRODUCTION.md
What didn’t work
Uncertainty-guided k-space acquisition tied with random selection (−0.06 ± 0.20 dB SEM, winning 31% of 32 slices). An encouraging n = 4 preliminary result turned out to be noise — reported as a null.
Reported as a null, not hidden — the same honesty policy applies to every number on this site.
Running it
No live reconstruction: a single image is ~1000 sampler steps. The site shows precomputed frames (zero-filled → reconstruction → ground truth) and the per-pixel uncertainty panel — the one output the original paper does not produce.