Generating Counterfactual Trajectories with Latent Diffusion Models for Concept Discovery

TL;DR

This paper addresses the challenge of trustworthy AI in high-stakes domains by proposing CDCT, a three-step framework that uses classifier-guided latent diffusion models to generate counterfactual trajectories, a VAE to disentangle semantic concepts, and a latent-space search to discover decision-relevant concepts. Applied to ISIC skin-lesion classification, CDCT uncovers both model biases and clinically meaningful biomarkers, while delivering counterfactuals with superior quality and up to 12x greater efficiency than prior diffusion-based methods. The approach enables unsupervised concept discovery in medical imaging, contributing to trust and knowledge expansion, with potential applicability to radiology and histology. Overall, CDCT advances explainable AI by combining latent-diffusion generation, disentangled latent representations, and data-driven concept identification to reveal how black-box classifiers make decisions in challenging domains.

Abstract

Trustworthiness is a major prerequisite for the safe application of opaque deep learning models in high-stakes domains like medicine. Understanding the decision-making process not only contributes to fostering trust but might also reveal previously unknown decision criteria of complex models that could advance the state of medical research. The discovery of decision-relevant concepts from black box models is a particularly challenging task. This study proposes Concept Discovery through Latent Diffusion-based Counterfactual Trajectories (CDCT), a novel three-step framework for concept discovery leveraging the superior image synthesis capabilities of diffusion models. In the first step, CDCT uses a Latent Diffusion Model (LDM) to generate a counterfactual trajectory dataset. This dataset is used to derive a disentangled representation of classification-relevant concepts using a Variational Autoencoder (VAE). Finally, a search algorithm is applied to identify relevant concepts in the disentangled latent space. The application of CDCT to a classifier trained on the largest public skin lesion dataset revealed not only the presence of several biases but also meaningful biomarkers. Moreover, the counterfactuals generated within CDCT show better FID scores than those produced by a previously established state-of-the-art method, while being 12 times more resource-efficient. Unsupervised concept discovery holds great potential for the application of trustworthy AI and the further development of human knowledge in various domains. CDCT represents a further step in this direction.

Figures (19)

• Figure 1: CDCT is a three-step concept discovery framework. An LDM with classifier guidance is used to generate a counterfactual trajectory dataset. A VAE is trained on this trajectory dataset to disentangle decision-relevant features. Finally, class-relevant dimensions are identified by manipulating the VAE's latent space and observing the target classifier's output.
• Figure 2: The counterfactual generation process starts by encoding the image $\mathcal{E}(x^F)$ and perturbing it to obtain $z_T$ (here $T=3$). The text encoder transforms the condition $(c)$ into an embedding $\tau_\theta(c)$. Both $z_T$ and $\tau_\theta(c)$ are fed to the diffusion model for denoising. At guided step $t$, we denoise the noisy latent $z_t$, $t$ times to produce the clean latent $v_t$, which is decoded into a clean image $\tilde{x_t}$ to calculate the gradient of loss $\mathcal{L_{\text{class}}}$ for updating $z_t$. We sample the previous less noisy latent $z_{t-1}$ from the estimated noise $\hat{\epsilon}$ and the updated noisy latent $\tilde{z_t}$. In the final guided step, $z_0$ is decoded to yield the counterfactual image $(x^{CF})$.
• Figure 3: An image of the Nevus class alongside its counterfactual images in all other target classes. The second row shows difference maps, providing an easy way to identify the areas of alteration between the original image and each counterfactual.
• Figure 4: Counterfactual trajectory for a Nevus with target class Melanoma. Along the process, the manifestation of darker, atypical pigment structures can be observed.
• Figure 5: Examples of discovered concepts using CDCT with the ISIC dataset. Each row shows two examples of a concept, where the original image, the reconstruction, and the manipulated reconstruction are aligned from left to right.