Publish the FLIP tutorial models

d4700c7 verified 7 days ago

9.95 kB

---
license: cc-by-sa-4.0
library_name: monai
tags:
- medical-imaging
- mri
- brain
- autoencoder
- latent-diffusion
- federated-learning
- nvflare
---
# Brain MRI latent diffusion: FLIP tutorial models

This repository holds the models trained with the two brain MRI tutorials of
[FLIP](https://github.com/londonaicentre/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
```text
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.
![The same four held-out patients reconstructed by v1, v2, v2.1, v3 and v3.1, axial, coronal and sagittal, with SSIM per reconstruction](figures/autoencoders_same_patients.png)
![The same patients reconstructed by v1.1, v2.1 and v3.1](figures/v1.1_v2.1_v3.1_same_patients.png)
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.

![The first three generated volumes per sequence from v3_ldm (E5), v2_ldm (E7), v3.1_ldm (E10) and v2.1_ldm (E11), beside three held-out patients](figures/diffusion_samples.png)

**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:

```python
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](https://arxiv.org/abs/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](https://arxiv.org/abs/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](https://arxiv.org/abs/1904.00625) (2019).

	---
	license: cc-by-sa-4.0
	library_name: monai
	tags:
	- medical-imaging
	- mri
	- brain
	- autoencoder
	- latent-diffusion
	- federated-learning
	- nvflare
	---

	# Brain MRI latent diffusion: FLIP tutorial models

	This repository holds the models trained with the two brain MRI tutorials of
	[FLIP](https://github.com/londonaicentre/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

	```text
	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.

	![The same four held-out patients reconstructed by v1, v2, v2.1, v3 and v3.1, axial, coronal and sagittal, with SSIM per reconstruction](figures/autoencoders_same_patients.png)

	![The same patients reconstructed by v1.1, v2.1 and v3.1](figures/v1.1_v2.1_v3.1_same_patients.png)

	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.

	![The first three generated volumes per sequence from v3_ldm (E5), v2_ldm (E7), v3.1_ldm (E10) and v2.1_ldm (E11), beside three held-out patients](figures/diffusion_samples.png)

	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:

	```python
	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](https://arxiv.org/abs/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](https://arxiv.org/abs/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](https://arxiv.org/abs/1904.00625) (2019).