Cycle Diffusion Model for Counterfactual Image Generation

TL;DR

This work addresses conditioning-faithful medical image synthesis by introducing Cycle Diffusion Model (CDM), which couples a cycle-consistent training objective with diffusion-based generation to produce accurate direct and counterfactual 3D brain MRIs conditioned on age and sex. CDM extends latent diffusion models with bidirectional (counterfactual and factual) generation and a cycle-regularization term, optimized via a composite loss that includes standard LDM denoising terms and a cycle-consistency penalty. Evaluated on a large, multi-study 3D brain MRI dataset, CDM outperforms baselines in conditioning accuracy (lower age MAE, higher sex accuracy), image quality (FID), and diversity (MS-SSIM), while also producing anatomically realistic age-related changes in counterfactuals. The approach offers potential for targeted data augmentation and disease progression modeling, though it incurs higher inference cost and requires careful consideration of distribution shifts and clinical validation for deployment.

Abstract

Deep generative models have demonstrated remarkable success in medical image synthesis. However, ensuring conditioning faithfulness and high-quality synthetic images for direct or counterfactual generation remains a challenge. In this work, we introduce a cycle training framework to fine-tune diffusion models for improved conditioning adherence and enhanced synthetic image realism. Our approach, Cycle Diffusion Model (CDM), enforces consistency between generated and original images by incorporating cycle constraints, enabling more reliable direct and counterfactual generation. Experiments on a combined 3D brain MRI dataset (from ABCD, HCP aging & young adults, ADNI, and PPMI) show that our method improves conditioning accuracy and enhances image quality as measured by FID and SSIM. The results suggest that the cycle strategy used in CDM can be an effective method for refining diffusion-based medical image generation, with applications in data augmentation, counterfactual, and disease progression modeling.

Figures (3)

• Figure 1: Graphical depiction of the proposed method. Our cycle diffusion model (CDM) performs a generative denoising process in the latent space in two directions. The counterfactual direction produces a counterfactual latent $\tilde{z}_0$ that is conditioned on the counterfactual condition $c'$. The factual direction produces a factual latent $\hat{z}_0$ that is conditioned on the original condition $c$. In addition to the denoising loss, a cycle-consistency loss is optimized which minimizes the distance between $z_0$ and $\hat{z_0}$.
• Figure 2: Representative real and synthetic samples for all models across the three planes are shown. The comparison includes the following methods, trained with the same experimental conditions as our method: VAE-GANlarsen2016autoencoding, $\alpha$-GAN kwon2019braingan, VQVAE-AR van2017neuraltudosiu2022autoregressive, and LDM rombach2022highresolutionpinaya2022brainimaginggenerationlatent. The selected subject is a 13-year-old male. We observe clearer ventricles and more detailed folding patterns on the cortex in CDM compared to baselines.
• Figure 3: Representative sample along three views of a 79-year-old female brain (left column) along with four counterfactuals and their difference maps with respect to the original sample: ($-60$, $-30$, $-10$, $+10$). Sex is unchanged. As age decreases, ventricle size shrinks (indicated by blue in difference maps) and the cortical surface appears thicker, reflecting reduced brain atrophy. In contrast, the older counterfactual shows enlarged ventricles (indicated by red in difference maps) and increased cortical thinning, consistent with age-related brain atrophy and neuro-degeneration gur1995sexXu112fjell2009minute. These gradual structural changes demonstrate the model's ability to realistically capture age-related brain morphology across different age intervals.