Federated training performs as well as centralized for 3D Latent Diffusion for Diagnosis-Conditioned Brain MRI Synthesis

Introduction

Deep learning for neuro-imaging is limited by data availability and distribution. Cohorts are often small, disease classes are imbalanced, and scanner and protocol differences between sites hurt generalization. Generative models offer one remedy, by synthesizing realistic, labeled scans for augmentation, for balancing rare classes, or for sharing data without exposing patient data. Denoising diffusion probabilistic models (DDPM) [1] now set the state of the art in image synthesis, and latent diffusion models (LDMs) [2] make them tractable for 3D medical volumes by running the diffusion process in the compressed latent space of an autoencoder. 3D LDMs have been used to generate high-resolution brain MRIs to complement real ones [3].

However, these models need large and diverse training data and brain MRI is held by separate hospitals and consortium, where data sharing is challenging due to privacy, patient consent, and data-governance constraints. Federated learning (FL) [4] allows training a deep learning model on multiple sites without sharing data. Each site is trained over its own private data, and only encrypted model parameters are exchanged and averaged (e.g., [5]). Diffusion models have been trained federatedly [6], and pretrained latent diffusion models have been adapted to heterogeneous clients in one-shot FL [7]. FL has been used for discriminative neuroimaging tasks as well [5]. Although latent diffusion has been proposed for privacy-preserving synthesis [8], training a full 3D brain MRI LDM federatedly remains much less explored.

We answer this question with a controlled comparison of the same diagnosis-conditioned 3D LDM, with identical architecture, losses and sampling where one is trained federatedly across five naturally-distributed data sites and another centrally on the pooled data, and both are evaluated on the same held-out scans with paired random seeds.

Data

We use T1-weighted MRI from the three phases of Alzheimer’s Disease Neuroimaging Initiative (ADNI) [9] (ADNI-1, ADNI-2/GO and ADNI-3) and the National Alzheimer’s Coordinating Center (NACC) [10], split into GE and Siemens/Philips scanner groups. Each of the five groups were represented as a federated site. Scans were reoriented, skull-stripped, bias-field corrected and affinely (9-DOF) registered to a common template at 2 mm isotropic resolution, padded to $96 \times112 \times96$ voxels and min–max normalized to $[0, 1]$. Each scan is labeled cognitively normal (CN) or Alzheimer’s disease (AD).

Table 1 summarizes the splits. Training, Validation and Test sets are subject disjoint. ADNI subjects contribute several longitudinal scans, while NACC is mostly one scan per subject.

Site	Train	Validation	Test
ADNI-1	807/741 (436)	76/113 (59)	88/127 (61)
ADNI-2/GO	1120/625 (593)	74/52 (48)	57/64 (39)
ADNI-3	847/188† (665)	–	–
NACC GE	978/761 (1530)	101/95 (165)	102/107 (180)
NACC Si/Ph*	891/363 (1048)	137/49 (157)	131/34 (142)
Total	4643/2678 (4028)	388/309 (418)	378/332 (417)

Methods

3D autoencoder

The autoencoder is a 3D AutoencoderKL from MONAI Generative Models [3, 11, 12] with channels $(64, 128, 256)$, two residual blocks per level, GroupNorm (16 groups) and self-attention at the lowest resolution. The encoder maps a volume $x$ to a Gaussian posterior with mean $\mu$ and standard deviation $\sigma$ over an $8 \times24 \times28 \times24$ latent space and the decoder reconstructs $x$ from a latent sample $z = \mu+ \sigma\odot\epsilon$. It is trained with

where $\mathcal{L}_{1}^{\mathrm{w}}$ is weighted L1 loss with brain voxels $(x > 0.05)$ weighted $\times10$, $\mathcal{L}_{\mathrm{perc}}$ is perceptual loss computed on 2D slices (20%) with a SqueezeNet backbone[13], and $\mathcal{L}_{\mathrm{adv}}$ is least-squares adversarial loss from a 3D PatchGAN discriminator, enabled after five warm-up epochs.

Diagnosis-conditioned latent diffusion

Following the original LDM design [2], the diffusion model is trained on sampled latents $z$ of the frozen autoencoder, multiplied by a scale factor $s = 1/\operatorname{std}(z)$ so that their variance matches the unit-variance diffusion noise. The denoiser is a 3D DiffusionModelUNet [14] with channels $(128, 256, 512)$, two residual blocks per level, and cross-attention at the two lower resolutions. The diagnosis $c \in\{\mathrm{CN}, \mathrm{AD}, \varnothing\}$ is mapped by a learnable embedding table to a 64-dimensional token consumed by cross-attention. We use a DDPM [1] forward process with $T = 1000$ steps in which noise is added gradually according to a variance schedule $\beta_1, \ldots, \beta_T$ (scaled-linear from $\beta_1 = 0.0015$ to $\beta_T = 0.0195$). $\bar{\alpha}_t = \prod_{i=1}^{t}(1-\beta_i)$ is the fraction of the signal that remains after $t$ steps. The UNet $\epsilon_\theta$ is trained to predict the added noise by minimizing

where $\epsilon\sim\mathcal{N}(0, I)$ is Gaussian noise of the latent’s size, $t$ is drawn uniformly from $\{1, \ldots, T\}$. For classifier-free guidance [15], the diagnosis is replaced by the null token $\varnothing$ with probability 0.15, so a single network learns both the conditional and the unconditional model. At sampling time, the two predictions are combined as

where the guidance scale $w = 1$ gives conditional and $w > 1$ strongly guided samples. We sample with Denoising Diffusion Implicit Model (DDIM) [16] (200 steps), decode $z_0/s$ with the autoencoder and clip the output to $[0,1]$.

Federated and centralized training

In both settings, the pipeline is trained in two consecutive stages. Stage 1 trains the autoencoder (Sec. 3.1) to compress each scan into a compact latent and reconstruct it and Stage 2 trains the diffusion UNet (Sec. 3.2).

Federated. Training runs on the SPHERE research testbed [17] with the MetisFL federated learning framework [18, 5] over five sites with the data in Table 1. Each site has two NVIDIA L40S GPUs. In each round every learner trains the community model for one local epoch on its own data, and the server aggregates the learners’ parameters by federated averaging, weighting each learner by its number of training scans. All five learners take part in every round, and each trains data-parallel on its two GPUs. Scans, latents and labels never leave a site. Sites share the model parameters each round (in stage 1, the autoencoder’s optimizer moments are also shared). During stage 1 the discriminator is kept local to each site and never shared, since only the autoencoder is needed downstream.

Centralized. The same models are trained on the pooled training set with distributed data-parallel training on 16 NVIDIA RTX A6000 GPUs (4 nodes). This is the upper bound that federated training is measured against and it requires all scans in one place.

Scale factor. The scale factor $s$ requires the standard deviation of the latents over all training scans. Each site encodes its scans with the frozen federated autoencoder and shares only three numbers: the count $N_k$, the sum $S_k = \sum z$ and the sum of squares $Q_k = \sum z^2$ over all latent values. The server computes $N = \sum_k N_k$, mean $= \sum_k S_k/N$, var $= \sum_k Q_k/N - \mathrm{mean}^2$ and $s = 1/\sqrt{\mathrm{var}}$, which equals the value computed on the pooled latents exactly. This yields $s = 0.9977$ for the federated autoencoder (per-site standard deviations 1.0007–1.0053) and $s = 1.0035$ for the centralized one.

Training details. All models are trained in full fp32 precision with Adam. Stage 1 uses a learning rate of $10^{-4}$ and a batch of 2 per GPU in both settings. In this stage, checkpoints are selected by validation PSNR within the first 500 epochs or rounds (centralized epoch 400, federated round 300). Stage 2 uses a batch of 10 per GPU, a learning rate of $5 \cdot10^{-5}$ with a linear warm-up, then constant, and gradient-norm clipping at 1.0. We evaluate both the centralized and federated models after 2000 epochs/rounds.

Evaluation

We generate 1,500 CN and 1,500 AD volumes at $w = 1$ (DDIM, 200 steps) using both the central and federated LDMs. Synthetic sets are compared with the pooled real test and validation scans (766 CN, 641 AD). To score of the generators, we also evaluate one random draw of 1,500 real training scans per class in exactly the same way.

• Conditioning fidelity (AD vs. CN). We train a 3D DenseNet-121 [19] AD vs CN classifier on the 7,321 labeled real training scans for 100 epochs with early stopping (patience=10) to prevent overfitting[20], using class-weighted cross-entropy and AdamW optimizer [21]. We report the AUC for each classifier.

• Diversity. Mean MS-SSIM [22] computed with MONAI’s MultiScaleSSIMMetric. We used 5 scales with the default weights, Gaussian kernel of size 4 and $\sigma= 1.5$, data range 1. MS-SSIM is computed for every pair of volumes within a set (each unordered pair $i < j$ once, no random subsampling), separately for each class, and averaged. Higher values mean more similar samples, i.e. lower diversity.

• Realism. Fréchet Inception Distance (FID) [23] on 2048-dimensional features of a 3D ResNet-50 pretrained on 23 medical datasets (MedicalNet [24]). The reference is always the pooled real test and validation scans. Lower FID means that the feature distribution of the synthetic volumes is closer to that of reference (real) scans. For computing FID, we used a brain mask for the real and synthetic images. Voxels above 0.01 after min–max scaling, of which only the largest 26-connected component is kept and its holes filled. Voxels outside the mask are set to 0 and those inside are z-scored, the intensity normalization MedicalNet was trained with.

Results

Reconstruction

Table 2 compares the selected autoencoders. Pooled over the four validation sites, the federated autoencoder performs similar to the centralized autoencoder with PSNR (26.09 vs. 26.30 dB) and equal SSIM. The ADNI sites are on par and the small gap is concentrated at the two NACC sites.

Site ($n$)	PSNR (dB) ↑ Fed	PSNR (dB) ↑ Cen	SSIM ↑ Fed	SSIM ↑ Cen	MAE ↓ Fed	MAE ↓ Cen
ADNI-1 (189)	25.51	25.48	0.981	0.978	0.0384	0.0386
ADNI-2/GO (126)	26.19	26.30	0.982	0.978	0.0352	0.0352
NACC GE (196)	26.19	26.55	0.973	0.976	0.0351	0.0341
NACC Si/Ph (186)	26.51	26.84	0.971	0.972	0.0341	0.0334
Pooled (697)	26.09	26.30	0.976	0.975	0.0357	0.0353

Generation

Table 3 compares federated and centralized training over the three metrics.

Class	Set	FID ($\times10^{-3}$)	MS-SSIM	AUC
CN	Cen	1.889	0.905	–
CN	Fed	1.254	0.918	–
CN	Real train	0.0284	0.896	–
CN	Real test+val	–	0.896	–
AD	Cen	1.878	0.900	–
AD	Fed	1.109	0.915	–
AD	Real train	0.0331	0.890	–
AD	Real test+val	–	0.890	–
Pooled	Cen	1.843	–	0.922
Pooled	Fed	1.143	–	0.934
Pooled	Real train	0.0156	–	–
Pooled	Real test+val	–	–	0.920

Conditioning fidelity is slightly better for federated samples. The real-data classifier recognizes the intended diagnosis of synthetic scans as reliably as that of real patients: AUC 0.934 for federated and 0.922 for centralized generated samples, vs. 0.920 for real test and validation scans. The federated–centralized AUC difference is small. Balanced accuracy (0.873 vs. 0.867), sensitivity (0.803 vs. 0.791) and specificity (0.943 in both) do not differ significantly. Both models produce CN samples that are more clearly CN than real CN scans (specificity 0.943 vs. 0.854) and AD samples that are somewhat less clearly AD than real AD scans (sensitivity 0.80 and 0.79 vs. 0.86).

Federated samples are more realistic. Federated samples are closer to real scans than centralized ones for both classes: (FID $\times10^{-3}$): 1.25 vs. 1.89 for CN, 1.11 vs. 1.88 for AD and 1.14 vs. 1.84 pooled. Both remain clearly distinguishable from real data, whose training draw scores FID 0.0156 · $10^{-3}$ against the held-out scans. In both models, 90–95% of FID comes from the mean-shift term $\lVert\mu_1 - \mu_2\rVert^2$, compared with 8–17% for the real training scans, whose small FID comes mostly from the covariance term. The synthetic sets thus vary about as much as real scans but are offset from them as a whole, which points to a systematic difference shared by both pipelines, plausibly the autoencoder’s reconstruction, rather than to the training paradigm.

Diversity is lower for federated samples. Centralized samples are nearly as diverse as real scans (MS-SSIM 0.905 and 0.900 for CN and AD, vs. 0.896 and 0.890 for both the real training draw and the real test and validation scans), whereas federated samples are more alike (0.918 and 0.915).

Discussion

Federated training is a viable alternative. Across reconstruction, conditioning fidelity and realism, the federated pipeline matches or exceeds the centralized reference with 99% of the autoencoder PSNR, an AD/CN classification of its samples at least as accurate as for centralized samples and for real patients, and lower FID for both diagnoses. However, federated samples cluster more tightly around typical anatomy and reproduce fewer of the atypical brains present in real data. A generator that concentrates on typical anatomy lands closer to the bulk of the real distribution, which lowers distribution distances dominated by the mean shift. A plausible cause is that averaging updates from heterogeneous sites pulls the model towards anatomy shared by all sites, the federated autoencoder may contribute as well. The diversity gap is moderate (0.918 vs. 0.905, with 0.896 for real data) and not a collapse with federated samples still carrying their diagnoses.

Privacy as the deciding factor. These small differences have to be weighed against what each setting requires. Centralized training needs every scan copied to one place, which multi-site clinical data rarely permits. In the federated setting, scans, latents and labels never leave their site. The only shared quantities are model parameters and, once, three summary statistics per site for the scale factor.

In conclusion, we trained the same diagnosis-conditioned 3D latent diffusion model for brain MRI federatedly across five sites and centrally on pooled data, and compared both on held-out test scans across three metrics: conditional fidelity, diversity, and realism. Federated training matched the centralized reference in autoencoder reconstruction (99% of PSNR) and in how reliably samples express the intended diagnosis, and produced samples closer to real scans in MedicalNet feature space, at the cost of moderately lower sample diversity. Since it achieves this without any scan leaving its site, federated training is a practical route to generative models built on multi-site clinical data that cannot be pooled.

Compliance with Ethical Standards

This retrospective analysis used de-identified data obtained from ADNI and NACC. All studies were approved by the institutional review boards of the participating sites, and participants provided written informed consent. The analysis used data under the applicable cohort data use agreements.

Acknowledgments

This work was supported in part by NIH grants RF1AG081571, U01AG068057, R01MH131806, and S10OD032285. We thank Dr. Jelena Mirkovic for providing GPU computing resources on the SPHERE testbed.