Pre to Post-Treatment Glioblastoma MRI Prediction using a Latent Diffusion Model

AI summary

Overview

Research area: Medical image generation / disease progression modeling — applying latent diffusion models to brain MRI in glioblastoma patients.

Technical level: Advanced (assumes familiarity with diffusion models, VAEs, latent spaces, and MRI imaging conventions, though the paper explains its pipeline step by step).

Scope: The paper proposes a 2D Latent Diffusion Model that generates a 4-month post-treatment T1-Gd MRI slice directly from a single pre-treatment T1-Gd MRI slice plus the tumor contour, optionally guided by predicted survival class.

What This Paper Is About

Glioblastoma patients respond very differently to the standard Stupp protocol, but the first visible treatment effect on MRI typically appears no earlier than two months after treatment begins — and the RANO response criteria are based on the 4-month post-treatment MRI. This paper asks whether a generative model can predict that 4-month post-treatment MRI slice from pre-treatment information alone, so clinicians could anticipate response (and tumor recurrence location) before treatment starts.

Key Contributions

1. A slice-to-slice formulation of treatment response prediction. Rather than modeling a time series of scans, the authors frame early visual Treatment Response Prediction as a translation from one pre-treatment MRI slice to one post-treatment MRI slice, which they argue better matches clinical practice.

2. A Latent Diffusion Model with concatenation-based conditioning. The model concatenates the pre-treatment MRI latent representation, a separately encoded Gross Tumor Volume (GTV) latent representation, and the noisy post-treatment latent, with a common MRI latent space and a dimensionality-matched GTV latent space.

3. Classifier-free guidance using survival information. A pre-trained classifier predicts survival class from the pre-treatment MRI, and that class embedding is injected into the UNet through cross-attention layers to steer generation. The authors compare a classifier predicted from images (ResNet-based, 83.78% accuracy) against ground-truth 2-class and 4-class survival labels, and against no class conditioning at all.

4. A trained, evaluated pipeline with released code. Trained on a local cohort of 140 GBM patients, evaluated with global and GTV-local image metrics plus a qualitative review by a medical physicist, with code released publicly.

Main Findings

• Dataset: 140 GBM patients treated at Centre François Baclesse between January 2018 and December 2023, with pre-treatment T1-Gd MRI, 4-month post-treatment T1-Gd MRI, expert-delineated GTV, and survival outcomes. Slices containing tumor yielded roughly 40 slices per patient and 6,059 slices total, at 256×256 pixels, split 80%/10%/10% at the patient level.

• VQ-VAE reconstruction quality: A shared latent space for pre- and post-treatment MRI achieved a PSNR of 42.08, an SSIM of 0.988, and an LPIPS of 0.008 on the test set.

• Class conditioning gave only modest gains: The LDM conditioned on the 2-class predicted classifier ("pred 2-class") achieved overall better results, while the LDM conditioned on 4-class ground-truth categories performed best on the GTV-local metrics. The authors conclude that the granularity of the outcome representation has limited impact. The exact per-model metric values appear in Table 1 of the paper and are not reproduced in the provided text.

• Qualitative expert scoring was comparable across models: A medical physicist scored one generated post-treatment T1-Gd slice per model for each of the 14 test patients on a 0–10 realism scale. Averages were 6.5 (no class, std 1.02, min 4, max 8), 6.0 (pred 2-class, std 0.80, min 4, max 7), 7.0 (2-class, std 0.80, min 6, max 8), and 7.0 (4-class, std 1.02, min 4, max 8).

• Characteristic failure mode: Across qualitative examples ordered by descending local SSIM (ranging from 0.71 down to 0.02), the model generally failed to predict T1-enhancement localized to the tumor area, but did succeed in predicting the necrotic area.

• Overall assessment: Predicting a post-treatment image from a single time-point representation of the disease "appeared to be a highly challenging task," and using a classifier had no significant impact on the post-treatment MRI prediction.

Methodology in Plain English

The authors built their system in two stages.

First, compression. They trained a Vector Quantized-Variational AutoEncoder (VQ-VAE) to squeeze MRI images into a small quantized latent space built from a learned codebook of embedding vectors, so that diffusion could run cheaply in that lower-dimensional space. One shared VQ-VAE was trained on both pre- and post-treatment MRI because the two are perceptually similar. A second VQ-VAE was trained separately on the tumor contour (GTV) because it looks very different from an MRI; its latent dimensionality was deliberately matched to the MRI latent so the two could be stacked together.

Second, generation. A Latent Diffusion Model based on Stable Diffusion was trained to predict the Gaussian noise added to the post-treatment latent. It receives three things stacked together — the pre-treatment MRI latent, the GTV latent, and the noisy post-treatment latent — so the model knows the general anatomy to reproduce and where the tumor is. On top of that, survival class is injected via a learned classification embedding applied at several UNet levels through cross-attention.

At inference, the model starts from pure noise and runs a reverse diffusion process using DDIM sampling with 200 steps. Each step makes two predictions: one conditioned on the survival class and one without conditioning. The final noise estimate is the conditioned prediction plus a guidance scale multiplier times the difference between the unconditioned and conditioned predictions — the standard classifier-free guidance trick, here with a guidance scale of 10.

Training details: PyTorch and Torch Lightning, two NVIDIA RTX A6000 GPUs (2 × 48 GB VRAM), batch size 30, AdamW optimizer with a learning rate of 2 × 10⁻⁶, 1,000 timesteps, 200 epochs with the best epoch selected on the validation set, and random horizontal flipping as the only augmentation.

Evaluation used PSNR, SSIM, LPIPS, and MSE between source and generated images, repeated locally within the GTV region because the goal is tumor evolution.

Why This Matters

Impact on research: The paper is an early attempt at a harder variant of disease progression modeling — predicting forward from a single time point rather than from a series of scans — and it reports honestly that the task is difficult. It also contributes a comparison of classifier-free guidance strategies (no class, predicted 2-class, ground-truth 2-class, ground-truth 4-class) showing that survival conditioning barely moves the needle, which is a useful negative result for the field. The released code and the local dataset description also give other groups a concrete starting point.

Real-world applications:

• Pre-treatment RANO scoring: Radiation oncologists apply RANO criteria using the 4-month post-treatment MRI; a predicted version could let them anticipate response before treatment begins.
• Radiotherapy planning: Spatial information about where recurrence is likely to appear could help optimize where radiation is delivered.
• Protocol personalization: Patients with very short versus much longer survival could in principle be identified earlier, supporting adaptation of therapeutic protocols.
• Trial stratification and patient communication: A predicted imaging trajectory could support earlier enrollment decisions and more concrete conversations about expected disease course.

Industry relevance: Latent diffusion pipelines conditioned on segmentation masks and tabular/class labels are directly reusable in medical imaging software, clinical trial tooling, and regulatory-facing synthetic-image workflows. The negative result on survival conditioning also matters to companies building AI response-prediction products, since it suggests image-level conditioning carries most of the signal.

Future Directions

1. Evaluation metrics aligned with clinical goals. The authors state plainly that current metrics lack correlation with the task objectives and that specific investigation is needed to identify an evaluation method matching clinical aims.

2. Adding 3D context. The current model is 2D slice-to-slice and lacks 3D context due to computational constraints; extending to volumetric generation is an explicit plan for further studies.

3. Improving prediction of T1-enhancement. The model generally failed to predict tumor-localized T1-enhancement while succeeding on the necrotic area — an open question is what architectural or conditioning change would fix this specific weakness.

4. Making survival guidance actually useful. Since classifier-free guidance with predicted or ground-truth survival classes produced no significant improvement, it remains open whether a better outcome representation, a stronger classifier, or a different conditioning mechanism could steer generation meaningfully.

Target Audience

Researchers and graduate students working on medical image generation, diffusion models, or disease progression modeling; medical physicists and radiation oncologists interested in computational prediction of treatment response in glioblastoma; and machine learning engineers in medical imaging who want a concrete, honest case study of a latent diffusion pipeline with mask and class conditioning, including what did not work.

Authors’ abstract

Glioblastoma (GBM) is an aggressive primary brain tumor with a median survival of approximately 15 months. In clinical practice, the Stupp protocol serves as the standard first-line treatment. However, patients exhibit highly heterogeneous therapeutic responses which required at least two months before first visual impact can be observed, typically with MRI. Early prediction treatment response is crucial for advancing personalized medicine. Disease Progression Modeling (DPM) aims to capture the trajectory of disease evolution, while Treatment Response Prediction (TRP) focuses on assessing the impact of therapeutic interventions. Whereas most TRP approaches primarly rely on timeseries data, we consider the problem of early visual TRP as a slice-to-slice translation model generating post-treatment MRI from a pre-treatment MRI, thus reflecting the tumor evolution. To address this problem we propose a Latent Diffusion Model with a concatenation-based conditioning from the pre-treatment MRI and the tumor localization, and a classifier-free guidance to enhance generation quality using survival information, in particular post-treatment tumor evolution. Our model were trained and tested on a local dataset consisting of 140 GBM patients collected at Centre François Baclesse. For each patient we collected pre and post T1-Gd MRI, tumor localization manually delineated in the pre-treatment MRI by medical experts, and survival information.