File size: 9,953 Bytes
d4700c7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | ---
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.
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:
```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).
|