Counterfactual Stress Testing for Image Classification Models

Abstract: Deep learning models in medical imaging often fail when deployed in new clinical environments due to distribution shifts in demographics, scanner hardware, or acquisition protocols. A central challenge is underspecification, where models with similar validation performance exhibit divergent real-world failure modes. Although stress testing has emerged as a tool to assess this, current methods typically rely on simple, uninformed perturbations (e.g., brightness or contrast changes), which fail to capture clinically realistic variation and can overestimate robustness. In this work, we introduce a counterfactual stress testing framework based on causal generative models that create realistic "what if" images by intervening on attributes such as scanner type and patient sex while preserving anatomical identity, enabling controlled and semantically meaningful evaluation under targeted distribution shifts. Across two imaging modalities (chest X-ray and mammography), three model architectures, and multiple shift scenarios, we show that counterfactual stress tests provide a substantially more accurate proxy for real out-of-distribution performance than classical perturbations, capturing the direction and relative magnitude of performance changes as well as model ranking. These results suggest that causal generative models can serve as practical simulators for robustness assessment, offering a more reliable basis for evaluating medical AI systems prior to deployment.

• The paper introduces a DSCM-based counterfactual stress testing framework that predicts out-of-distribution performance more accurately than classical perturbation tests.
• It employs hierarchical β-Variational Autoencoders to generate high-fidelity counterfactual images by isolating clinically relevant variables like scanner type and sex.
• Empirical results on PadChest and EMBED datasets show high Pearson correlations (up to r=0.95) with real OOD performance, underscoring its value for pre-deployment robustness evaluation.

Counterfactual Stress Testing for Robustness Evaluation in Medical Imaging

Introduction

This paper addresses the challenge of robust model deployment for medical image classification under distribution shift, focusing on the well-documented issue that high in-distribution performance is not indicative of reliability under deployment conditions in new clinical environments. The central contributions are a causal generative model-based counterfactual stress testing framework for image classifiers and an extensive empirical evaluation comparing its predictive value against classical perturbation stress tests. The work systematically investigates failures associated with underspecification—where models sharing similar validation accuracy exhibit divergent failure modes on out-of-distribution (OOD) data—by leveraging Deep Structural Causal Models (DSCMs) to simulate targeted, clinically realistic distribution shifts.

Methodology

The proposed framework employs DSCMs realized as hierarchical $\beta$-Variational Autoencoders (HVAEs) to synthesize high-fidelity counterfactual images. These generative models are conditioned on specific, causally relevant parent variables (notably scanner type and biological sex), while preserving patient-specific anatomy. This enables the generation of "counterfactual twins" that isolate the effects of targeted covariates—a substantial methodological advance over conventional data augmentation or pixel-level perturbations. Causal disentanglement is promoted via a staged KL regularization schedule, and model stability is enhanced by gradient clipping and EMA-based weight averaging.

For empirical assessment, the PadChest (chest X-ray, $n\approx160,000$) and EMBED (mammography, $n\approx290,000$) datasets are utilized, with rigorous splits by patient ID and attribute-matched OOD test sets. Three standard architectures—ResNet-50, DenseNet-121, and ViT-B/16—are evaluated across tasks including pneumonia detection and BI-RADS breast density classification. Stress tests are performed by generating counterfactual images varying either scanner type or demographic attribute, and performance deltas ($\Delta$AP or $\Delta$AUC) between these and the original IID test sets are compared to real OOD performance as a proxy for true deployment robustness. Classical perturbation stress tests (brightness, contrast, sharpness, gamma, blur) are included as baselines.

Experimental Results and Claims

Robustness Under Covariate Shift

Across both imaging modalities and all architectures, the counterfactual stress testing approach dramatically outperforms classical perturbation tests in predicting real OOD performance shifts. Notably, counterfactual tests achieve a Pearson correlation of $r=0.93$ (PadChest, scanner) and $r=0.85$ (PadChest, sex) with true OOD $\Delta$AP, compared to maximal classical perturbation $r=0.47$ and frequently much lower or negative values. For composite distribution shifts (joint scanner and sex interventions), correlation increases further ($r=0.95$), indicating robust alignment with real-world deployment conditions. Mean absolute error (MAE) between predicted and real performance shifts is consistently minimized by the counterfactual approach, and this holds across attribute, model architecture, and dataset. In particular, the method also captures both the direction and magnitude of performance change under OOD shifts—something classical approaches routinely fail to achieve.

In the EMBED mammography experiments, counterfactual stress testing again outperforms perturbation, achieving Pearson $n\approx160,000$0 (counterfactual) versus statistically insignificant correlations for all classical methods. Despite some underestimation of the magnitude of scanner-induced degradation, the method reliably models the sign and comparative impact of shifts across multiple scanners.

Key Claims

• Counterfactual stress tests yield a substantially more accurate proxy for real OOD performance than classical perturbation-based tests across both modalities and all tested model architectures.
• Classical stress tests consistently underestimate (or misalign with) the impact of clinically meaningful distribution shifts, while counterfactual testing closely approximates both the direction and relative magnitude of degradation.
• Composite shift scenarios confirm the scalability of the approach to more complex, clinically realistic covariate interactions.

Theoretical and Practical Implications

The work significantly advances the methodology for pre-deployment robustness assessment in high-stakes medical AI. By simulating attribute-specific interventions on a generative model's latent space, the method can expose underspecified failure modes that are undetectable by perturbation-based techniques or in-distribution validation metrics. The causal modeling ensures that generated counterfactuals maintain anatomical verisimilitude, which is crucial for downstream interpretability and clinical acceptance. This directly addresses gaps in model ascertainment relevant to regulatory and audit frameworks seeking fairness and reliability in AI-assisted diagnosis.

Practically, the framework enables scenario-based robustness audits and has clear applicability to pre-clinical validation, regulatory reporting, and ongoing monitoring in deployed systems. Because the approach isolates causal factors, it can also inform targeted data augmentation or retraining strategies, and may be further extensible to simulate intersectional demographic and technical shifts.

Limitations and Future Directions

Two primary limitations are acknowledged. First, the method presupposes knowledge (or reliable estimation) of the underlying causal graph and identifiability of the relevant counterfactuals within the generative model. These are strong assumptions, especially in multi-factorial settings where unobserved confounding and causal ambiguity are likely. Second, the evaluation relies on simulated counterfactuals, whose fidelity—while empirically validated—is ultimately limited by the capacity and disentanglement of the DSCM. There remains a need for quantitative and expert-informed realism metrics for synthetic images.

Future research directions include development of data-driven, structure-learning DSCMs (e.g., with approaches akin to DECI), quantitative evaluation metrics for counterfactual realism, extension to three-dimensional or multimodal imaging, multi-attribute interventions, and integration of human-in-the-loop verification by clinical experts.

Conclusion

This study presents a causally informed, generative framework for counterfactual stress testing in medical image classification. The empirical evidence supports the conclusion that model robustness under controlled, attribute-specific counterfactual interventions is a significantly closer proxy to real-world OOD generalization than perturbation-based stress testing. The approach has meaningful implications for the safety, fairness, and auditability of medical AI, and lays the groundwork for more sophisticated robustness certification methods based on causal modeling and counterfactual reasoning.