aicentreflip/tutorials-brain-mri-latent-diffusion

Brain MRI latent diffusion: FLIP tutorial models

This repository holds the models trained with the two brain MRI tutorials of FLIP, the Federated Learning Interoperability Platform, in fl-tutorials/nvflare/image_generation/:

• autoencoder/: 3-D autoencoders (MONAI AutoencoderKL, 21.0 M parameters) that compress a brain MR volume into a small latent.
• diffusion/: 3-D latent diffusion models (MONAI DiffusionModelUNet, 49.3 M parameters) that generate in an autoencoder's latent space, conditioned on the MR sequence (FLAIR, T1w, T1Gd, T2w).

Every model was trained federated, with NVIDIA FLARE simulating two sites on one GPU. v3.1 is the autoencoder to use and v3.1_ldm the diffusion model trained on it; the tutorials fetch them. The other folders are the steps that led there, published so every score below can be checked.

These are tutorial models, trained on a small public research dataset. They are not for clinical use.

Contents

autoencoder/<version>/         v1  v1.1  v2  v2.1  v3  v3.1
    autoencoder.safetensors    the weights, float32 as trained
    autoencoder_config.yaml    the architecture, and the fingerprint the diffusion job checks
    scores_heldout.json        SSIM per held-out volume
    provenance.json            commit, install line, launch command, versions, hardware, wall time
    code/                      the app files exactly as the run executed them
diffusion/<autoencoder>_ldm/   v2_ldm  v2.1_ldm  v3_ldm  v3.1_ldm, each named by its frozen autoencoder
    diffusion_model.safetensors
    scores.json                realism, coverage, sequence probe
    provenance.json            as above, plus the autoencoder's folder and fingerprint
    code/
splits/                        the 321 training studies and the 40 held-out patients
evaluation/                    the scripts behind every score and figure on this card
figures/

Autoencoders

All six were scored on the same 40 held-out patients (160 volumes, four sequences each). Each reconstruction is mapped back onto the original 1 mm scan and compared there with SSIM inside the brain, so models with different input geometries are scored alike.

Folder	Run	Geometry	Training data	Rounds × epochs	Learning rate	Latent channels	SSIM at 1 mm	Wall time
v1	E1	squashed to 96³	100 studies	3 × 150	constant 1e-4	3	0.683	7 h 09 min
v1.1	E4	squashed to 96³	321 studies	10 × 25	constant 1e-4	3	0.723	11 h 47 min
v2	E2	2 mm, 96×112×80	321 studies	5 × 25	1e-4 → 1e-5 each round	3	0.699	6 h 49 min
v2.1	E9	2 mm, 96×112×80	321 studies	10 × 25	constant 1e-4	3	0.776	12 h 25 min
v3	E3	2 mm, 96×112×80	321 studies	5 × 25	1e-4 → 1e-5 each round	8	0.765	6 h 14 min
v3.1	E8	2 mm, 96×112×80	321 studies	10 × 25	constant 1e-4	8	0.823	12 h 28 min

v1 is the tutorial's original recipe, which squashes each 240×240×155 scan into a 96³ cube. The 2 mm geometry resamples to 2 mm voxels and pads or crops to 96×112×80 instead, so the brain keeps its shape. v1.1, v2.1 and v3.1 train in bfloat16, as the tutorial does now; v1, v2 and v3 used float16. Per sequence, v3.1 scores FLAIR 0.786, T1w 0.844, T1Gd 0.832 and T2w 0.830, and it beats v2.1 and v3 on all 160 held-out volumes.

Training: L1 reconstruction, a MedicalNet ResNet-50 perceptual loss, a very small KL weight (1e-8) and a least-squares adversarial loss from a 3-D patch discriminator (not published: only the autoencoder is needed afterwards).

Latent diffusion models

Each was trained for 5 rounds × 60 epochs on the same 321 studies, with the learning rate falling from 1e-4 to 1e-6 over the 300 epochs, conditioned on the sequence. It then generated 40 volumes per sequence (160 in all), scored against the 40 held-out patients on MedicalNet ResNet-50 features.

Folder	Run	Frozen autoencoder	Latent	MMD², mean ↓	Precision ↑	Recall ↑	Sequence probe ↑	Malformed by eye ↓	Wall time
v2_ldm	E7	v2	3 × 24 × 28 × 20	0.111	0.91	0.77	0.506	7 / 160	5 h 15 min
v3_ldm	E5	v3	8 × 24 × 28 × 20	0.070	0.91	0.81	0.625	5 / 160	5 h 15 min
v2.1_ldm	E11	v2.1	3 × 24 × 28 × 20	0.130	0.88	0.78	0.463	14 (+10 borderline) / 160	5 h 16 min
v3.1_ldm	E10	v3.1	8 × 24 × 28 × 20	0.053	0.96	0.78	0.650	4 (+3 borderline) / 160	5 h 35 min
held-out scans					0.88	0.93	0.894
v3.1 reconstructions				−0.019		0.92	0.825

• MMD² (Gaussian kernel): 0 means indistinguishable from the held-out scans.
• Precision is the share of generated volumes that land among real training scans, and recall the share of real scans the generated volumes reach (k-nearest neighbours, k = 3).
• Sequence probe: a logistic regression trained to tell the four sequences apart on real scans (0.894 on held-out ones) reads each generated volume; the score is the share read as the sequence that was asked for.

Limitations. No model here generates tumours, although every training scan has one: tumours differ in place, size and shape from patient to patient, and the diffusion models learn the brain without them. As a result they return T1w-like volumes when asked for T1Gd, which differs from T1w mainly in the enhancing tumour rim. v3.1_ldm's 40 volumes per sequence read as:

Asked for	FLAIR	T1w	T1Gd	T2w
FLAIR	23	14	2	1
T1w	1	39	0	0
T1Gd	0	34	6	0
T2w	2	2	0	36

A few volumes in every run are malformed (blurred or fragmented brains); the count is in the table.

Use

The tutorials fetch these files themselves. From fl-tutorials/nvflare/image_generation/brain_mri_ldm_stage2_diffusion/, make prepare-checkpoint downloads autoencoder/v3.1, checks its hash, and writes the two files the diffusion job needs.

To load a model directly:

import yaml
from monai.networks.nets import AutoencoderKL
from safetensors.torch import load_file

config = yaml.safe_load(open("autoencoder/v3.1/autoencoder_config.yaml"))
autoencoder = AutoencoderKL(**config["autoencoder"])
weights = load_file("autoencoder/v3.1/autoencoder.safetensors")
autoencoder.load_state_dict({k.removeprefix("autoencoder."): v for k, v in weights.items()})

The diffusion UNet's architecture and its LATENT_SCALE_FACTOR are in its code/config.json; its tensors carry a diffusion_model. prefix the same way. The inputs must go through the run's own code/transforms.py: z-scored per volume, and resampled to the geometry in the table.

Reproducing a result

1. Install flip-utils from the folder's commit, with the line in its provenance.json. The runs used branch commits that are in no PyPI release.
2. Check out FLIP at the same commit and copy the folder's code/ over the tutorial's app_files/. The runs predate the tutorials' move to image_generation/, so provenance.json gives the directory they ran from.
3. Download the data with make -C fl-tutorials download-brain-mri-msd-raw and download-brain-mri-data, and train on splits/dataframe_train.csv. v1 trained on studies BRATS_001 to BRATS_100 instead.
4. Run the launch command, then the matching script from evaluation/ on the result.

Training on a GPU is not bit-for-bit repeatable (cuDNN kernels, random augmentation, sampled latents), so a rerun should land within run-to-run noise of these scores rather than on them. To check a number without training, run the evaluation script on the published weights.

All runs used one NVIDIA GeForce RTX 5090, PyTorch 2.13.0, MONAI 1.6.0 and NVIDIA FLARE 2.9.0.

Data and licence

Trained on Task01 (BrainTumour) of the Medical Segmentation Decathlon: 484 multiparametric MR studies of gliomas from the BraTS 2016 and 2017 challenges, each with FLAIR, T1w, T1Gd and T2w. Some patients appear in several studies, so the data were split by patient: 40 patients are held out whole for scoring (splits/heldout_cases.txt), and training keeps every other patient's studies except exact re-uploads, 321 studies (splits/dataframe_train.csv). The dataset is licensed CC BY-SA 4.0, and so are these models.

The perceptual loss and the diffusion scores use MedicalNet (Chen, Ma and Zheng, 2019, arXiv:1904.00625; MIT licence), whose weights are not included.

References

• M. Antonelli et al., "The Medical Segmentation Decathlon", Nature Communications 13, 4128 (2022).
• A. L. Simpson et al., "A large annotated medical image dataset for the development and evaluation of segmentation algorithms", arXiv:1902.09063 (2019).
• W. H. L. Pinaya et al., "Brain imaging generation with latent diffusion models", DGM4MICCAI (2022).
• S. Chen, K. Ma and Y. Zheng, "Med3D: Transfer learning for 3D medical image analysis", arXiv:1904.00625 (2019).