Multi-site MRI harmonization with StarGAN
Mapping T1 scans from three cohorts into an ADNI reference domain, with a fourth site the model never sees.
- Scans
- 9,740
- Subjects
- 3,450
- Eval configs
- 12 per task
Overview
A domain-conditional 3D StarGAN (six-block residual generator, PatchGAN discriminator with an auxiliary domain classifier) maps T1 MRI from NACC, OASIS3 and NIMHANS into an ADNI reference domain at 96×128×128.
A data-scaling study ran five 100-epoch trainings spanning ADNI 1k→2k scans, NACC 500→1k, a full-ADNI cohort of 7,939 scans and a size-matched configuration: 3,450 subjects and 9,740 scans across four sites, with NIMHANS held out entirely as a zero-shot target.
Challenges
The training loop had no way to choose a checkpoint, so I built one: rank all 100 epochs on adversarial equilibrium, reconstruction and domain-classification losses, then score the shortlist empirically on Wasserstein distance to the real ADNI intensity distribution, change magnitude, and SSIM, the last specifically to rule out generators that had collapsed to an identity map.
Harmonized outputs came back near-black. The cause was InstanceNorm applying training running statistics at inference; per-volume statistics restored 0.96–0.98 voxel correlation with the input.
The downstream evaluation models were Keras with TF-Addons, a dependency that had broken. Porting them to PyTorch meant layer-by-layer numerical verification of activations, optimizer updates, early-stopping behaviour and metric accumulation.
Decisions
The calls that shaped this project, the alternatives they beat, and what each one cost.
Build a checkpoint-selection methodology
over take the final epoch
WhyGAN loss does not track output quality monotonically, and the training loop offered no selection at all. Ranking epochs on equilibrium and reconstruction, then scoring the shortlist on Wasserstein distance and SSIM, is what separates a working generator from one that learned the identity map.
Trade-offA second evaluation pass over every shortlisted checkpoint.
Per-volume InstanceNorm statistics at inference
over the running statistics learned during training
WhyRunning statistics produced near-black outputs. Computing normalization from the scan being harmonized rather than from the training distribution restored 0.96–0.98 voxel correlation with the input.
Trade-offInference is no longer independent of the volume it is given.
Hold NIMHANS out entirely as a zero-shot site
over train across all four sites
WhyA harmonizer that only works on sites it has seen is no use to the next hospital. Keeping one site unseen measures generalization instead of memorization.
Trade-offGives up that site's data during training.
Port the downstream models from Keras to PyTorch
over pin the broken TF-Addons dependency
WhyThe evaluation models are the measuring instrument. A dependency that no longer installs makes every downstream number unreproducible, so the port was verified layer by layer against the original.
Trade-offWeeks spent on infrastructure that produces no new result by itself.
Compute fidelity over every axial slice
over a single mid-slice, as is conventional
WhyA generator can look excellent on the central slice and fall apart at the volume edges. Per-slice metrics, broken out per diagnosis and per site, show where it actually fails.
Trade-offFar more computation across 12 configurations.
Outcomes
Image fidelity is measured with PSNR, SSIM, mutual information, FID and MMD computed over every axial slice rather than one mid-slice, reported per diagnosis and per site.
Downstream impact is measured directly: CN versus Dementia classification and brain-age regression (MAE, MAPE, RMSE, R², Pearson r) across 12 harmonized and unharmonized train/test configurations per task, on 500-scan subject-disjoint test sets per site.
The pipeline now runs on the full CN/MCI/Dementia cohort rather than CN-only.