Scanner-agnostic 3D brain MRI harmonization
Reimplementing a 3D GAN that separates anatomy from scanner style, and finding out where it breaks.
- Scans
- 3,816
- Resolution
- 182×218×182
- GPU budget
- 17.5 GB
Overview
AGUIT-3D (Caldera et al., 2026) is a 3D GAN that separates brain anatomy from scanner style. No public implementation exists, so this is a reimplementation in PyTorch from the paper and its supplement: anatomy and style encoders, an AdaIN generator, multi-scale image and latent critics, and SSIM-based structure and luminance losses.
Training runs on a 3,816-scan multi-site cohort of 1 mm T1 MRIs from ADNI, NACC, OASIS3 and NIMHANS, at the full 182×218×182 resolution, using mixed precision and gradient checkpointing to fit inside 17.5 GB.
Challenges
Evaluation had to be model-agnostic so the reproduction and its successor could be judged on the same footing: image-similarity metrics (PSNR, SSIM, mutual information), distribution distances to the reference site (FID, MMD), and two downstream tasks (CN versus dementia classification and brain-age regression) on subject-disjoint test sets across three external cohorts.
The model looked like it was working by one measure and failing by another, so the diagnosis came from targeted GPU experiments rather than from reading loss curves.
Decisions
The calls that shaped this project, the alternatives they beat, and what each one cost.
Reimplement from the paper and supplement
over wait for code, or use an off-the-shelf harmonizer
WhyNo public implementation of AGUIT-3D exists. Building it from the paper is the only way to benchmark the approach, and it forces every architectural assumption to be explicit rather than inherited.
Trade-offAny gap between the reproduction and the authors' intent is indistinguishable from a finding.
Train at the full 182×218×182 resolution
over downsample to fit the GPU
WhyHarmonization artifacts live in fine intensity structure, so downsampling hides exactly what the method is supposed to fix. Mixed precision and gradient checkpointing bought full resolution inside a 17.5 GB budget.
Trade-offSlower steps and a smaller batch size than a downsampled run.
Probe whether the style code carries site identity
over trust that the architecture disentangles
WhyDisentanglement is an assumption, not a guarantee. Testing whether the style encoder can identify the source site at all turned a vague sense that results looked wrong into a specific failure: it was at chance.
Judge on distribution distance and downstream tasks
over report intensity alignment
WhyThe histograms matched, with distance falling up to 55%, while FID rose 1.5–2× and brain-age predictions shifted 3–4 years. A single-axis metric would have reported this as a success.
Trade-offRequires training and maintaining downstream models alongside the generator.
Outcomes
The disentanglement does not hold. The style encoder could not identify the source site, and the latent critic stayed at chance. The model was not learning the separation it is built around.
Harmonized scans matched ADNI intensity histograms, with distance down up to 55%, but their FID to ADNI rose 1.5–2× and an ADNI-trained model rated them 3–4 years older. Matching a histogram is not the same as matching a distribution.
Next: replacing AdaIN with StyleGAN2-style weight demodulation to remove normalization artifacts, benchmarked against this reproduction through the same pipeline.