Latent Drifting in Diffusion Models for Counterfactual Medical Image Synthesis

TL;DR

This paper introduces Latent Drift, a method to improve diffusion models for medical image synthesis, enabling better counterfactual image generation despite limited data and domain shifts.

Contribution

The paper proposes Latent Drift, a novel approach that enhances diffusion models' ability to generate counterfactual medical images under data scarcity and domain shift conditions.

Findings

Significant performance improvements in counterfactual medical image generation.

Effective across multiple datasets and fine-tuning schemes.

Facilitates investigation of clinical factors like age, gender, and diseases.

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