AI Alignment in Medical Imaging: Unveiling Hidden Biases Through Counterfactual Analysis

TL;DR

The paper tackles bias in medical-imaging ML by proposing a counterfactual-invariance (CI) based statistical test to quantify dependence on sensitive attributes, circumventing the need for counterfactual data. It develops a practical CIT-LR algorithm that leverages a disentangled latent diffusion model (CLDM) and Mutual Information Neural Estimation (MINE) to learn CI-relevant representations and test for invariance via a paired $t$-test. Across synthetic and real datasets (CheXpert, MIMIC-CXR), the method shows strong alignment with counterfactual fairness (ECA) and outperforms traditional DP/EO baselines in detecting true invariances and biases. The approach provides a robust, causal tool for evaluating generalization across demographic groups in AI-driven healthcare, while acknowledging limitations in latent-specification assumptions and diffusion variance that warrant further refinement and interdisciplinary validation.

Abstract

Machine learning (ML) systems for medical imaging have demonstrated remarkable diagnostic capabilities, but their susceptibility to biases poses significant risks, since biases may negatively impact generalization performance. In this paper, we introduce a novel statistical framework to evaluate the dependency of medical imaging ML models on sensitive attributes, such as demographics. Our method leverages the concept of counterfactual invariance, measuring the extent to which a model's predictions remain unchanged under hypothetical changes to sensitive attributes. We present a practical algorithm that combines conditional latent diffusion models with statistical hypothesis testing to identify and quantify such biases without requiring direct access to counterfactual data. Through experiments on synthetic datasets and large-scale real-world medical imaging datasets, including \textsc{cheXpert} and MIMIC-CXR, we demonstrate that our approach aligns closely with counterfactual fairness principles and outperforms standard baselines. This work provides a robust tool to ensure that ML diagnostic systems generalize well, e.g., across demographic groups, offering a critical step towards AI safety in healthcare. Code: https://github.com/Neferpitou3871/AI-Alignment-Medical-Imaging.

Figures (14)

• Figure 1: (a) Visualization of the conditional latent diffusion model used in Alg. \ref{['alg']}. (b) Examples of generated images using the CLDM of Alg. \ref{['alg']} under different conditioning labels. The bottom row highlights differences between the original image and the generated images. Larger counterfactual images are presented in App. \ref{['app:large_image_generation']}.
• Figure 2: The relationship between log-transformed $p$-values and the ECA tested on synthetic datasets for $100$ base classifiers (see App. \ref{['app:synthetic_classifiers']}). Column names refer to the different synthetic datasets (see App. \ref{['app:synthetic_data_description']}). Each point represents the tested result of a spECAfic classifier. A straight line is fitted to highlight the Pearson correlation, and the parallel painted region shows the standard deviations of the $p$-values. The red dashed horizontal line represents a significance threshold at $\log (p)=\log (0.05)$. Our testing method aligns more consistently with the ground truth then baselines.
• Figure 3: The relationship between log-transformed $p$-values and the ECA is examined using the cheXpert dataset (DBLP:conf/aaai/IrvinRKYCCMHBSS19, App. \ref{['app:synthetic_data_description']}) for $100$ base classifiers (App. \ref{['app:chexpert_classifiers']}). The column labels correspond to the various target diseases of the dataset. Each point represents the outcome of testing a spECAfic classifier. A straight line is fitted to illustrate the Pearson correlation, with a shaded area indicating the standard deviation of the $p$-values. A red dashed horizontal line marks the significance threshold at $\log (p) = \log(0.05)$. Our testing method aligns more consistently with the ground truth than baselines.
• Figure 4: The relationship between log-transformed $p$-values and the ECA is examined using the MIMIC-CXR dataset (mimic-cxrmimic-cxr-jpg, App. \ref{['app:synthetic_data_description']}) for $100$ base classifiers (App. \ref{['app:chexpert_classifiers']}). The column labels correspond to the various target diseases of the dataset. Each point represents the outcome of testing a spECAfic classifier. A straight line is fitted to illustrate the Pearson correlation, with a shaded area indicating the standard deviation of the $p$-values. A red dashed horizontal line marks the significance threshold at $\log (p) = \log(0.05)$. Our testing method aligns more consistently with the ground truth then baselines.
• Figure 5: Causal structure for the DGP of the CLDM as in Fig. \ref{['fig:generative_model']}.

Theorems & Definitions (8)

• Definition 4.1: Counterfactual Invariance, following DBLP:conf/nips/KusnerLRS17quizan.2024
• Theorem 4.2
• Definition A.1: Structural Causal Model (SCM), Definition 2.1 by 10.1214/21-AOS2064
• Definition D.1: Valid Adjustment Set
• Lemma D.2
• Lemma D.3
• proof
• proof : Proof of Theorem \ref{['lemma:cond_variance']}