Latent Drifting in Diffusion Models for Counterfactual Medical Image Synthesis

Yousef Yeganeh 1,2*, Azade Farshad 1,2*†, Ioannis Charisiadis 1‡, Marta Hasny 1‡, Martin Hartenberger 1‡, Björn Ommer 2,3, Nassir Navab 1,2, Ehsan Adeli 4

1Technical University of Munich, Munich, Germany    2Munich Center for Machine Learning, Munich, Germany   3CompVis @ LMU Munich, Munich, Germany   4Stanford University, Stanford, CA, USA

y.yeganeh@tum.de, azade.farshad@tum.de,

*,‡Equal Contribution, †Project Lead

(🔆 Highlight Paper) 📌 Conference on Computer Vision and Pattern Recognition (CVPR), 2025

Collaborating Institutions

Paper Code (Coming Soon) BibTeX
Latent Drifting Teaser Image
Image Generation w. and w/o. LD during fine-tuning. Examples generated from left to right using Textual Inversion.
Latent Drifting Teaser Image
Generated counterfactual samples on CheXpert using Pix2Pix Zero + LD.

Abstract

Scaling by training on large datasets has been shown to enhance the quality and fidelity of image generation and manipulation with diffusion models; however, such large datasets are not always accessible in medical imaging due to cost and privacy issues, which contradicts one of the main applications of such models to produce synthetic samples where real data is scarce. Also, fine-tuning on pre-trained general models has been a challenge due to the distribution shift between the medical domain and the pre-trained models. Here, we propose Latent Drift (LD) for diffusion models that can be adopted for any fine-tuning method to mitigate the issues faced by the distribution shift or employed in inference time as a condition. Latent Drifting enables diffusion models to be conditioned for medical images fitted for the complex task of counterfactual image generation, which is crucial to investigate how parameters such as gender, age, and adding or removing diseases in a patient would alter the medical images. We evaluate our method on three public longitudinal benchmark datasets of brain MRI and chest X-rays for counterfactual image generation. Our results demonstrate significant performance gains in various scenarios when combined with different fine-tuning schemes.

Methodology

Latent Drifting Method Overview
Figure 2: Latent Drifting (LD) enables diffusion models to generate high-quality medical images by controlling the trade-off between diversity and condition adherence through the latent drift parameter (δ).

Our proposed method, Latent Drifting (LD), is formulated as a min-max optimization problem aiming to match the learned distribution of pre-trained models to a new distribution represented by finite accessible samples. The latent space adds to the traditional conditions (e.g., text or image), and its underlying distribution functions as an additional hyperparameter and conditioning factor.

Key Components

δ = -0.1 δ = -0.05 δ = 0 δ = 0.05 δ = 0.1
Musk δ=-0.1 Musk δ=-0.05 Musk δ=0 Musk δ=0.05 Musk δ=0.1
Obama δ=-0.1 Obama δ=-0.05 Obama δ=0 Obama δ=0.05 Obama δ=0.1
Figure 3: Samples generated with identical sampled noise and varying latent drift (δ ∈ [-0.1,0.1]) at inference, with different text prompts. Top: "Elon Musk on a mountain", Bottom: "Barack Obama on a plane".

Applications

Results

Alzheimer's Disease Transformation

AD Source
Source (AD)
CN Counterfactual
LD (Ours)
Difference
Difference

Healthy to Alzheimer's Disease

CN Source
Source (Healthy)
AD Counterfactual
LD (Ours)
Difference
Difference

Brain Aging

Young Source
Source (Young)
Old Target
Target (Old)
LD Result
LD (Ours)

Chest X-ray Disease Generation

Real Pneumonia
Real Sample
Without LD
Without LD
With LD
With LD (Ours)

We evaluate our method on three public longitudinal benchmark datasets:

Our experiments demonstrate that Latent Drifting significantly improves the performance of diffusion models for medical image generation and manipulation across various scenarios:

Conclusion

We presented Latent Drifting (LD), a novel approach for enhancing diffusion models for medical image generation and manipulation. Our method addresses the challenges of distribution shift between natural and medical images, enabling high-quality counterfactual medical image synthesis. By formulating conditioning as a counterfactual explanation optimization problem, LD provides a more restrictive way of conditioning suited for medical applications.

The key advantages of our approach include:

Future work will explore the application of Latent Drifting to other medical imaging modalities and investigate its potential for clinical decision support through counterfactual reasoning.

BibTeX

        
@inproceedings{yeganeh2025latent,
  title={Latent Drifting in Diffusion Models for Counterfactual Medical Image Synthesis},
  author={Yeganeh, Yousef and Farshad, Azade and Charisiadis, Ioannis and Hasny, Marta and Hartenberger, Martin and Ommer, Björn and Navab, Nassir and Adeli, Ehsan},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2025}
}
        
      

Acknowledgement

This work and collaboration were partially supported by the Munich Center for Machine Learning, BaCaTeC, and the Stanford HAI Hoffman-Yee Award.