Technical Summary: AI Alignment in Medical Imaging via Counterfactual Analysis

Problem Statement

Machine learning (ML) systems in medical imaging, while demonstrating remarkable diagnostic capabilities, are susceptible to learning spurious correlations with sensitive attributes (e.g., self-reported race, gender). These biases can lead to disparities in clinical outcomes, such as misdiagnosis or undertreatment for specific demographic groups. A critical challenge is that ML models can identify demographic factors from medical images—often with accuracy surpassing human capability—by relying on proxies or surrogate covariates correlated with these attributes. Existing fairness evaluation methods, such as Demographic Parity (DP) and Equality of Opportunity (EO), rely on statistical associations and often fail to detect hidden biases or misinterpret legitimate clinical variations as unfairness. Furthermore, direct observation of counterfactual scenarios (e.g., the same patient with a different race) is impossible in reality, making it difficult to rigorously test whether a model's predictions are invariant to changes in sensitive attributes.

Methodology

The paper introduces a novel statistical framework to evaluate the dependency of medical imaging ML models on sensitive attributes using the concept of Counterfactual Invariance (CI). The core hypothesis is that a model is fair if its prediction remains unchanged when a sensitive attribute is hypothetically altered, provided all other clinical information remains constant.

The proposed methodology, termed Conditional Independence Test via Latent Representations (CIT-LR), consists of three main components:

1. Disentangled Conditional Latent Diffusion Models (DCLDM):

   • To simulate counterfactual data, the authors employ a generative framework based on latent diffusion models.
   • The architecture includes an encoder-decoder pair and a diffusion process operating in a lower-dimensional latent space.
   • A key innovation is a nonlinear transformation ($\phi_{\lambda,\delta}$) applied to the latent representation. This transformation is trained to minimize the mutual information between the latent code and the sensitive attribute, effectively "disentangling" the attribute from the clinical content (anatomy, pathology).
   • The model generates high-quality counterfactual images by conditioning the diffusion process on a new attribute value (e.g., changing "Black" to "White") while preserving the disentangled clinical features.

2. Theoretical Framework and Algorithm:

   • The approach is grounded in Structural Causal Models (SCMs). A model $f$ is defined as counterfactually invariant with respect to attribute $A$ if $E[\hat{y} | do(A=a), Z=z] = E[\hat{y} | do(A=a'), Z=z]$, where $Z$ represents attribute-invariant factors.
   • Algorithm 1 operationalizes this by:
     • Training a DCLDM to learn disentangled representations.
     • Generating counterfactual images for test samples across different attribute values.
     • Computing a statistic comparing the model's prediction on the generated image against the expected value over all possible attribute interventions.
     • Applying a statistical hypothesis test (paired t-test for binary outputs, Kolmogorov–Smirnov test for continuous outputs) to determine if the model's predictions are statistically invariant to the attribute intervention.

3. Evaluation Metrics:

   • Rejection Rate (Statistical Power): The probability of correctly rejecting the null hypothesis (that the model is invariant) for biased models.
   • False Alarm Rate: The probability of incorrectly rejecting invariance for fair models.
   • Expected Counterfactual Accuracy (ECA): A ground-truth metric used on synthetic datasets to measure the extent of counterfactual invariance.

Key Contributions

• Novel Statistical Framework: The paper proposes the first practical algorithm to test counterfactual invariance in pretrained medical imaging models without requiring access to real counterfactual data.
• Disentangled Generation: It introduces a DCLDM architecture specifically designed to generate realistic counterfactual medical images where sensitive attributes are altered while disease-relevant information is preserved.
• Rigorous Testing Pipeline: The method combines generative modeling with formal statistical hypothesis testing, moving beyond simple association-based metrics (DP/EO) to a causal evaluation of model behavior.
• Comprehensive Validation: The framework is validated on both synthetic benchmarks (where ground-truth invariance is known) and two large real-world chest X-ray datasets (CheXpert and MIMIC-CXR).

Experimental Results

The framework was evaluated against existing baselines, including Demographic Parity (DP), Equality of Opportunity (EO), and other counterfactual generation methods (non-disentangled CLDM and ImageCFGen).

• Synthetic Datasets:

  • CIT-LR correctly identified 96.1% of biased models (ECA < 0.8).
  • It produced a false alarm rate of 22.0% on near-invariant models (ECA $\ge$ 0.8).
  • The method significantly outperformed DP and EO, which showed non-monotonic behavior and lower rejection rates for biased models. It also outperformed non-disentangled baselines, highlighting the importance of the disentanglement step.

• Real-World Datasets (CheXpert and MIMIC-CXR):

  • Across five diagnostic tasks (Cardiomegaly, Edema, Consolidation, Atelectasis, Pleural Effusion) and three racial groups, CIT-LR detected bias at average rates of 96.3% (CheXpert) and 95.7% (MIMIC-CXR) for strongly biased models.
  • For the least biased models, the false alarm rate dropped to 15.3% (CheXpert) and 14.7% (MIMIC-CXR).
  • Baselines (DP, EO) exhibited irregular performance, often failing to detect bias in clearly biased models or flagging fair models inconsistently.

• Generative Quality and Validity:

  • Effectiveness Accuracy: The DCLDM achieved high accuracy (88.7%–93.7%) in generating images with the target racial attribute.
  • Realism: Low Fréchet Inception Distance (FID) scores (21.79 for CheXpert, 25.89 for MIMIC-CXR) indicated high visual fidelity.
  • Artifact Detection: Classifiers trained to distinguish real from synthetic images performed at chance levels (AUC $\approx$ 0.5), confirming the absence of detectable digital artifacts.
  • Signal Removal: Post-conversion, classifiers could no longer distinguish between racial groups (AUC dropped from ~0.97 to ~0.56), confirming the successful removal of demographic signals.

Significance and Claims

The paper claims that this work represents a significant step toward building trustworthy and unbiased AI systems in healthcare by aligning model evaluation with counterfactual fairness principles.

• Beyond Association: The authors argue that traditional metrics like DP and EO are insufficient because they rely on statistical associations that may conflate legitimate clinical prevalence differences with unfair bias. Their method directly tests the causal dependency of predictions on sensitive attributes.
• Practical Applicability: By providing a post-hoc testing procedure that does not require retraining models or access to unobservable counterfactual data, the framework offers a practical tool for auditing existing medical AI systems.
• Alignment with Causal Principles: The results demonstrate that the proposed framework adheres closely to the theoretical notion of counterfactual invariance, offering a more robust measure of fairness than current baselines.

The authors remain modest regarding limitations, acknowledging that the validity of the test relies on the DCLDM's ability to perfectly disentangle attributes and that the use of proxy metrics (ECA) in real-world settings introduces some uncertainty due to the inability to observe true counterfactuals. They also note that the current study focuses on single sensitive attributes and that future work must address joint attributes and unobserved confounding.