WorksComputer Science
Counterfactual Marginalisation: Framework for Evaluating Robustness to Nuisance Variables
Abstract
Machine learning models can achieve strong test performance while relying on demographic or acquisition-related shortcuts. We propose counterfactual (CF) marginalisation as a test-time evaluation procedure for assessing robustness of classification models to such variables. Given a CF image generator, we intervene on nuisance parent variables such as age or sex, generate CF versions of each test image, and average predictions over a target intervention distribution. This produces intervention-aware predictions that marginalise demographic effects while preserving patient-specific latent information. We use these predictions to define metrics for CF risk, calibration, stability and worst-case sensitivity. We demonstrate this framework's utility for quantitative robustness evaluation.
Introduction
Deep learning models have achieved strong performance in medical image analysis, but high accuracy on held-out data does not guarantee clinically appropriate decision rules. In practice, models may rely on demographic, acquisition-related, or site-specific shortcuts rather than pathology-relevant evidence [12]. Causal modelling provides a principled way to study these issues by asking, for example, how a prediction would change if a patient’s demographic attributes were intervened upon while other patient-specific factors were held fixed [21]. Recent work has used Structural Causal Model (SCM)-based generative models to produce medical image CFs for explanation, auditing and augmentation [20, 22, 4, 9]. Here, instead, we use CFs at test time to define evaluation metrics for medical image classifiers. Given an SCM-based CF generator, we intervene on nuisance parent variables such as age or sex, and generate CF versions of each test image [9]. We then develop a range of metrics that use these CFs to measure the sensitivity, stability and performance of a predictor. This setting is closely related to work on fairness and robustness in medical imaging [12, 16], where prior studies have shown that models can exhibit demographic disparities and can recover sensitive attributes directly from images [28, 2]. Our contributions are as follows:
We formulate counterfactual marginalisation over variables as a test-time evaluation procedure for medical image classifiers.
We introduce intervention-aware evaluation metrics, including counterfactual marginal risk, interventional expected risk, counterfactual stability, and worst-case counterfactual risk. This extends counterfactuals from explanation and augmentation to principled robustness evaluation.
We demonstrate that our approach exposes biases in predictors better than traditional metrics and also enables measuring sensitivity to nuisance variables without requiring disease labels.
Background
Point-estimate evaluation can miss brittle model behaviour: small, plausible changes to inputs or logits may alter predictions. This concern echoes work on flat minima, where models in wider low-loss regions tend to generalise better than sharp solutions [8, 5]. Ensembles and Bayesian model averaging address a related issue by marginalising over parameters to improve robustness and uncertainty estimates [13, ?, 26], while test-time augmentation similarly averages predictions over transformed images [1, 14]. Our counterfactual marginalisation can be viewed as a causal analogue: instead of averaging over ad-hoc image transformations, we average over SCM-defined interventions on parent variables such as age or sex. This in turn motivates the design of metrics that assess not only average performance, which often is insufficient for describing clinical utility [19], but also metrics that measure sensitivity, calibration, and worst-case behaviour across clinically meaningful perturbations.
Counterfactual Marginalisation for Evaluation
Let denote the predicted probability of disease for image and the corresponding ground-truth label. We partition the parent variables as , where is the subset of parent variables to be marginalised, e.g. age and sex, and denotes the remaining parent variables that are kept fixed. Given an observed sample , we generate counterfactual images by intervening on while preserving the inferred exogenous noise for that individual through the standard abduction–action–prediction procedure. This yields a set of counterfactual images , where denotes the abducted exogenous variables for sample . To remove the influence of the nuisance demographic variables, we define a target distribution over interventions. This may be chosen as the empirical test-set distribution, a uniform distribution over demographic strata, or an external population distribution. We then define the counterfactually marginalised prediction for sample as

Figure 1. Overview of the proposed Counterfactual Marginalisation evaluation framework
In practice, this expectation is approximated by Monte Carlo sampling:
Counterfactual Risk
Using (1), we define the counterfactual marginal risk and interventional expected risk as
where and is binary cross-entropy or another appropriate loss. The former evaluates predictive performance after marginalising out the nuisance demographic variables, while the latter averages the loss over counterfactual interventions before aggregation. To measure the extent to which the predictor is vulnerable to the most adverse demographic interventions, we consider a tail-risk measure based on the conditional value at risk (CVaR). For a chosen tail level , the counterfactual tail risk is
where is the expected loss within the worst -fraction of interventions under . Let denote the risk of the worst-case intervention.
Counterfactual Calibration
While and measure predictive accuracy under counterfactual marginalisation, they do not directly assess whether the resulting probabilities are well calibrated. To handle this, we report the counterfactual Brier score and counterfactual expected calibration error (ECE):
where denote bins of samples partitioned according to their marginalised predicted probabilities . These measure whether the probabilities obtained after marginalisation remain clinically interpretable as calibrated risk estimates.
Counterfactual Stability
We define a counterfactual stability score based on the variation of the model output across interventions:
A low value of indicates that the predicted probability is stable under changes in the nuisance demographic variables. For threshold-based clinical decision making, we additionally define a counterfactual flip rate
where is a fixed decision threshold. This metric captures how often a prediction changes solely because of an intervention on the nuisance parent variables. We also measure how often the binary decision made on the observed image differs from the decision made using the counterfactually marginalised prediction. This observed–marginalised disagreement rate is defined as
This metric quantifies the extent to which conventional observational evaluation would lead to a different clinical decision than intervention-aware evaluation. For tasks in which the predicted score is used for prioritisation or triage, it is important that the relative ranking of patients remains stable after marginalisation [3, 17]. Let denote the rank of the observed prediction across the test set, and let denote the rank of the marginalised prediction . We quantify counterfactual rank stability using the Spearman rank correlation measure:
Theoretical Advantages of Counterfactual Evaluation
As well as the clear advantage of being able to generate samples from underrepresented classes, we also have that for quantities that are identified observationally, such as the mean prediction within a demographic subgroup, the CF estimator achieves lower variance by exploiting the full test set rather than the subgroup alone. Let denote a particular attribute value (e.g. ), the number of test samples in that subgroup, and the total test size. The standard observational estimator of is
The CF estimator generates a counterfactual image at for every test patient and averages: . Under correct SCM specification, is unbiased for . Its variance decomposes via the law of total variance, conditioning on the abducted exogenous noise :
The first term is the variance of the posterior mean (signal); the second is the expected posterior variance (SCM generation noise ). The exogeneous noise accounts for all unobserved background variables and individual traits, serving as a sufficient statistic for individual-level causal queries. Then, by the Rao-Blackwell theorem, , so
Proposition 1 (Variance reduction). Let . The CF estimator satisfies
whenever . The advantage grows as the subgroup becomes rarer: as , the bound reduces to .
This is directly analogous to the efficiency gain of a within-subjects design over a between-subjects design in experimental settings [23].
Experiments
Datasets
To demonstrate the utility of our metrics, we use two large scale chest x-ray datasets, CheXpert [10] and MIMIC-CXR [11]. We perform binary classification of Pleural Effusion (PE), existence or not. To ensure a controlled evaluation, we balance the test split so that male and female patients appear in equal numbers, and disease prevalence is equalised across sex and 3 age groups (, -, ). The final balanced test sets for CheXpert (C) and MIMIC (M) contain 5,180 and 28,189 studies respectively. We use these two datasets to run three experimental settings: (1) Train and test the PE classifier on C; (2) Train and test on M; (3) Train on C and test on with natural demographic distributions. For all 3 settings we train the CF generator on C.
Predictors and Training Settings
The causal graph encodes the following parents of the image: sex, age, race, lung size, and pleural effusion. All experiments vary the sex attribute () and age (8 uniformly-spaced values from 20 to 90 years), yielding counterfactual images per patient. The marginalisation distribution is uniform over the intervention values in all experiments. The baseline predictor (basic) is trained on the full training set without any demographic manipulation. To study how well our method assesses bias, we create a first set of biased predictors by adding a fixed scalar offset to the baseline logit for female patients at test time: , where is the sex label. Because the offset is a monotone transformation applied uniformly within each sex group, within-group rankings are preserved and hence stratified AUC is invariant to by construction. We then train three sex-biased predictors on a corrupted dataset in which Male+disease and Female+healthy samples are dropped with probabilities 0.8, 0.9, and 0.95 respectively, inducing a spurious correlation between sex and PE. We also train predictors with a continuous age bias that induces a correlation between older age and positive PE. For each training sample with normalised age and binary PE label , the concordance score is defined as , and a sample is removed from training with probability , where controls the strength of the induced correlation.
Results
tbl:3a shows that in the CheXpert experiment (Setting 1), our CF metrics discern the biases in the various predictors better than the standard risk (cross entropy) calculated only on original data, which sometimes decreases when bias increases (e.g. Logit+0.5 Logit+1 ). To summarise comparisons, we define the Correct Transition Count (CTC) as the number of transitions with increasing predictor bias (out of 9 total: 3 for logit bias, Basic Logit+0.5 +1 +2, and similarly 3 for sex and age bias) where the risk also increases. Higher is better and our counterfactual metrics outperform the standard risk estimation . All results shown are averages of 3 repetitions per experiment with varying seeds. Similar results were found for setting 2 and 3 (See Appendix A for details).
Moreover, biased predictors should be worse calibrated [18]. Figure 2 compares the original ECE [7] and marginalised ECE ([]{#eq:7}) for each predictor, where we find that scores of the latter increase more consistently with increasing predictor bias. This suggests that CF marginalised ECE can better estimate model calibration.

Figure 2. Model calibration: Marginalised ECE increases more consistently with predictor bias than the original ECE, suggesting it is a better measure of calibration.
| Predictor | |||||||
| Basic | 0.444 | 0.501 | 0.594 | 1.002 | 1.326 | 1.549 | 1.852 |
| Logit +0.5 | 0.445 | 0.612 | 0.724 | 1.027 | 1.346 | 1.597 | 1.845 |
| Logit +1 | 0.443 | 0.629 | 0.771 | 1.068 | 1.424 | 1.699 | 1.962 |
| Logit +2 | 0.502 | 0.677 | 0.919 | 1.072 | 1.431 | 1.716 | 1.997 |
| Sex () | 0.453 | 0.610 | 0.687 | 1.192 | 1.618 | 1.991 | 1.390 |
| Sex () | 0.452 | 0.614 | 0.699 | 1.195 | 1.612 | 1.968 | 2.338 |
| Sex () | 0.443 | 0.614 | 0.699 | 1.203 | 1.622 | 1.982 | 2.361 |
| Age () | 0.494 | 0.613 | 0.750 | 1.126 | 1.511 | 1.817 | 2.111 |
| Age () | 0.509 | 0.612 | 0.754 | 1.225 | 1.654 | 2.007 | 2.340 |
| Age () | 0.512 | 0.613 | 0.758 | 1.518 | 2.046 | 2.495 | 2.893 |
| CTC () | 6.33 | 8.33 | 8.66 | 8.33 | 8.00 | 8.00 | 7.33 |
Table 1. Marginal risk decomposition – PE classifier is trained and tested on CheXpert.
Finally, we assess our stability metrics. When ground truth disease labels are available, bias is often estimated by calculating the performance gap between the best and worst performing data stratification for a variable of interest. Therefore, separately for each attribute (sex, race, age), we calculate the worst-case gap in AUC () between relevant data stratifications (e.g. male/female for sex) for the 1 non-biased (basic) predictor and each of the 3 sex-biased and 3 age-biased predictors. We repeat for 3 random seeds and thus obtain 21 values of per attribute (sex, race, age). We then study how well our stability metrics correlate with and Figure 3 demonstrates strong correlation, implying our metrics are useful bias estimators. Table 5c further shows that values of our stability metrics tend to follow changes in model bias. These results suggests that our metrics can be used as bias estimators without requiring ground truth classification labels, replacing the use of stratified performance gap analysis, which requires ground truth labels, when these are not available.

Figure 3. Our stability metrics correlate strongly with AUC gaps, suggesting they are a useful proxy when labels are unavailable at test time.
Table 5: Prediction stability under marginalisation at threshold (uniform prior). = variance of marginalised predictions; = label-flip rate at ; = mean absolute difference between original and marginalised scores; = Spearman rank correlation; = Kendall ; = mean absolute rank change.
| Predictor | ||||||
| Basic | 0.0396 | 0.247 | 0.199 | 0.691 | 0.476 | 913.5 |
| Logit +0.5 | 0.0530 | 0.271 | 0.216 | 0.653 | 0.471 | 918.4 |
| Logit +1 | 0.0574 | 0.289 | 0.235 | 0.653 | 0.469 | 928.5 |
| Logit +2 | 0.0685 | 0.335 | 0.291 | 0.641 | 0.466 | 934.2 |
| Sex () | 0.0426 | 0.265 | 0.218 | 0.653 | 0.476 | 920.6 |
| Sex () | 0.0501 | 0.256 | 0.208 | 0.645 | 0.479 | 922.4 |
| Sex () | 0.0489 | 0.263 | 0.222 | 0.644 | 0.464 | 932.5 |
| Age () | 0.0566 | 0.267 | 0.203 | 0.652 | 0.477 | 913.3 |
| Age () | 0.0567 | 0.270 | 0.205 | 0.642 | 0.466 | 937.8 |
| Age () | 0.0567 | 0.279 | 0.208 | 0.589 | 0.416 | 1050.2 |
Table 5. (a) Trained and tested on CheXpert ().
| Predictor | ||||
| Basic | 0.0198 | 0.143 | 0.069 | 0.920 |
| Logit +0.5 | 0.0212 | 0.144 | 0.064 | 0.918 |
| Logit +1 | 0.0359 | 0.188 | 0.115 | 0.907 |
| Logit +2 | 0.0502 | 0.243 | 0.173 | 0.847 |
| Sex () | 0.0270 | 0.146 | 0.068 | 0.916 |
| Sex () | 0.0305 | 0.156 | 0.079 | 0.906 |
| Sex () | 0.0284 | 0.158 | 0.082 | 0.903 |
| Age () | 0.0263 | 0.167 | 0.077 | 0.917 |
| Age () | 0.0308 | 0.170 | 0.079 | 0.916 |
| Age () | 0.0314 | 0.174 | 0.096 | 0.901 |
Table 2. Stability under marginalisation for cross-dataset experiment. As desired, increase in bias leads in most cases to increase of , , , and decrease of . This consistent relation indicate that the metrics can be used to estimate levels of bias.
Conclusion
This work introduced counterfactual marginalisation as a test-time framework for evaluating medical image classifiers under controlled interventions on nuisance parent variables. By averaging predictions over SCM-generated counterfactual variants, the framework provides intervention-aware measures of risk, calibration, stability, tail sensitivity, and ranking robustness. Results on CheXpert and MIMIC demonstrate that counterfactual evaluation can expose sensitivity to demographic interventions even when standard point-estimate metrics appear stable. A limitation is that the metrics depend on the realism and causal validity of the counterfactual generator, so they should be used alongside generator validation and conventional external evaluation, rather than as a replacement. Future work should study sensitivity to causal graph specification, intervention distributions, generator uncertainty, and extensions to multi-label settings.
Full Results
Risk
| Predictor | ||||||||
| (a) Trained and tested on CheXpert (). | ||||||||
| Basic | 0.444 | 0.501 | 0.594 | 1.002 | 1.326 | 1.549 | 1.852 | |
| Logit +0.5 | 0.445 | 0.612 | 0.724 | 1.027 | 1.346 | 1.597 | 1.845 | |
| Logit +1 | 0.443 | 0.629 | 0.771 | 1.068 | 1.424 | 1.699 | 1.962 | |
| Logit +2 | 0.502 | 0.677 | 0.919 | 1.072 | 1.431 | 1.716 | 1.997 | |
| Sex () | 0.453 | 0.610 | 0.687 | 1.192 | 1.618 | 1.991 | 1.390 | |
| Sex () | 0.452 | 0.614 | 0.699 | 1.195 | 1.612 | 1.968 | 2.338 | |
| Sex () | 0.443 | 0.614 | 0.699 | 1.203 | 1.622 | 1.982 | 2.361 | |
| Age () | 0.494 | 0.613 | 0.750 | 1.126 | 1.511 | 1.817 | 2.111 | |
| Age () | 0.509 | 0.612 | 0.754 | 1.225 | 1.654 | 2.007 | 2.340 | |
| Age () | 0.512 | 0.613 | 0.758 | 1.518 | 2.046 | 2.495 | 2.893 | |
| CTC () | 6.33 | 8.33 | 8.66 | 8.33 | 8.00 | 8.00 | 7.33 | |
| (b) Trained and tested on MIMIC (). | ||||||||
| Basic | 0.335 | 0.425 | 0.536 | 0.790 | 1.031 | 1.190 | 1.283 | |
| Logit +0.5 | 0.342 | 0.454 | 0.556 | 0.827 | 1.084 | 1.252 | 1.351 | |
| Logit +1 | 0.359 | 0.469 | 0.590 | 0.897 | 1.161 | 1.370 | 1.494 | |
| Logit +2 | 0.424 | 0.511 | 0.694 | 0.880 | 1.181 | 1.377 | 1.505 | |
| Sex () | 0.375 | 0.455 | 0.570 | 1.055 | 1.290 | 1.459 | 1.563 | |
| Sex () | 0.371 | 0.458 | 0.583 | 1.215 | 1.427 | 1.580 | 1.681 | |
| Sex () | 0.384 | 0.459 | 0.581 | 1.232 | 1.499 | 1.724 | 1.864 | |
| Age () | 0.335 | 0.425 | 0.536 | 0.857 | 1.160 | 1.360 | 1.549 | |
| Age () | 0.351 | 0.565 | 0.734 | 0.887 | 1.142 | 1.370 | 1.501 | |
| Age () | 0.366 | 0.610 | 0.771 | 0.880 | 1.191 | 1.411 | 1.505 | |
| CTC () | 7.00 | 8.33 | 8.00 | 8.00 | 7.66 | 7.00 | 7.33 | |
| (c) Trained on CheXpert, tested on MIMIC (). | ||||||||
| Basic | 0.397 | 0.466 | 0.549 | 0.769 | 0.977 | 1.121 | 1.210 | |
| Logit +0.5 | 0.402 | 0.486 | 0.566 | 0.822 | 1.055 | 1.223 | 1.328 | |
| Logit +1 | 0.422 | 0.495 | 0.602 | 0.913 | 1.176 | 1.379 | 1.506 | |
| Logit +2 | 0.508 | 0.536 | 0.722 | 1.177 | 1.506 | 1.778 | 1.945 | |
| Sex () | 0.418 | 0.497 | 0.558 | 0.761 | 0.941 | 1.064 | 1.140 | |
| Sex () | 0.395 | 0.499 | 0.565 | 0.793 | 1.023 | 1.183 | 1.281 | |
| Sex () | 0.437 | 0.530 | 0.597 | 0.829 | 1.042 | 1.193 | 1.287 | |
| Age () | 0.416 | 0.490 | 0.568 | 0.897 | 1.056 | 1.143 | 1.464 | |
| Age () | 0.422 | 0.511 | 0.573 | 0.809 | 1.015 | 1.166 | 1.563 | |
| Age () | 0.402 | 0.516 | 0.614 | 0.827 | 1.062 | 1.233 | 1.642 | |
| CTC () | 6.66 | 8.00 | 8.33 | 8.33 | 7.66 | 7.66 | 7.33 |
Table 3. Marginal risk decomposition (uniform prior). = risk of original predictor; = risk of the complete-marginalisation predictor; = individual-environment risk; = conditional value-at-risk at level ; = worst-case risk; CTC () = correct transition count.
Calibration
| Predictor | Marginalised Brier | Marginalised ECE | Marginalised MCE | Original Brier | Original ECE | Original MCE |
| (a) Trained and tested on CheXpert (). | ||||||
| Predictor | Brier | ECE | MCE | Brier | ECE | MCE |
| Basic | 0.202 | 0.083 | 0.171 | 0.143 | 0.030 | 0.080 |
| Logit +0.5 | 0.208 | 0.125 | 0.219 | 0.143 | 0.019 | 0.077 |
| Logit +1 | 0.216 | 0.153 | 0.240 | 0.149 | 0.043 | 0.121 |
| Logit +2 | 0.237 | 0.211 | 0.306 | 0.180 | 0.116 | 0.228 |
| Sex () | 0.218 | 0.100 | 0.184 | 0.146 | 0.017 | 0.046 |
| Sex () | 0.223 | 0.104 | 0.232 | 0.145 | 0.042 | 0.119 |
| Sex () | 0.224 | 0.111 | 0.233 | 0.142 | 0.017 | 0.078 |
| Age () | 0.216 | 0.091 | 0.164 | 0.158 | 0.058 | 0.104 |
| Age () | 0.216 | 0.099 | 0.159 | 0.162 | 0.057 | 0.089 |
| Age () | 0.228 | 0.103 | 0.171 | 0.165 | 0.056 | 0.079 |
| (b) Trained and tested on MIMIC (). | ||||||
| Predictor | Brier | ECE | MCE | Brier | ECE | MCE |
| Basic | 0.139 | 0.088 | 0.244 | 0.102 | 0.012 | 0.067 |
| Logit +0.5 | 0.143 | 0.111 | 0.284 | 0.104 | 0.029 | 0.093 |
| Logit +1 | 0.150 | 0.136 | 0.323 | 0.109 | 0.050 | 0.146 |
| Logit +2 | 0.167 | 0.179 | 0.378 | 0.128 | 0.091 | 0.212 |
| Sex () | 0.152 | 0.126 | 0.299 | 0.193 | 0.133 | 0.262 |
| Sex () | 0.186 | 0.128 | 0.303 | 0.212 | 0.154 | 0.271 |
| Sex () | 0.207 | 0.135 | 0.296 | 0.241 | 0.181 | 0.328 |
| Age () | 0.141 | 0.094 | 0.251 | 0.113 | 0.039 | 0.149 |
| Age () | 0.150 | 0.105 | 0.252 | 0.114 | 0.049 | 0.188 |
| Age () | 0.154 | 0.138 | 0.259 | 0.118 | 0.042 | 0.073 |
| (c) Trained on CheXpert and tested on MIMIC (). | ||||||
| Predictor | Brier | ECE | MCE | Brier | ECE | MCE |
| Basic | 0.155 | 0.059 | 0.103 | 0.125 | 0.056 | 0.147 |
| Logit +0.5 | 0.157 | 0.074 | 0.161 | 0.126 | 0.030 | 0.077 |
| Logit +1 | 0.161 | 0.096 | 0.206 | 0.134 | 0.019 | 0.034 |
| Logit +2 | 0.179 | 0.150 | 0.278 | 0.164 | 0.075 | 0.146 |
| Sex () | 0.156 | 0.061 | 0.116 | 0.132 | 0.075 | 0.163 |
| Sex () | 0.169 | 0.061 | 0.139 | 0.125 | 0.065 | 0.182 |
| Sex () | 0.172 | 0.076 | 0.132 | 0.140 | 0.072 | 0.187 |
| Age () | 0.159 | 0.063 | 0.141 | 0.131 | 0.040 | 0.075 |
| Age () | 0.167 | 0.064 | 0.137 | 0.135 | 0.028 | 0.083 |
| Age () | 0.170 | 0.096 | 0.163 | 0.127 | 0.015 | 0.048 |
Table 4. Calibration before and after marginalisation (10 bins). “Marginalised” = marginalised predictions; “Original” = original predictions. Brier = Brier score; ECE = expected calibration error; MCE = maximum calibration error.
Stability
| Predictor | ||||||
| Basic | 0.0358 | 0.148 | 0.069 | 0.932 | 0.783 | 2143 |
| Logit +0.5 | 0.0367 | 0.148 | 0.071 | 0.930 | 0.782 | 2059 |
| Logit +1 | 0.0385 | 0.153 | 0.080 | 0.909 | 0.760 | 2242 |
| Logit +2 | 0.0433 | 0.170 | 0.108 | 0.877 | 0.745 | 2731 |
| Sex () | 0.0406 | 0.158 | 0.102 | 0.919 | 0.758 | 2041 |
| Sex () | 0.0497 | 0.196 | 0.137 | 0.897 | 0.713 | 3013 |
| Sex () | 0.0514 | 0.191 | 0.139 | 0.865 | 0.671 | 3386 |
| Age () | 0.0412 | 0.155 | 0.087 | 0.928 | 0.773 | 2244 |
| Age () | 0.0477 | 0.151 | 0.089 | 0.917 | 0.770 | 2285 |
| Age () | 0.0487 | 0.168 | 0.097 | 0.888 | 0.746 | 2330 |
Table 5b. (b) Trained and tested on MIMIC ().
| Predictor | ||||||
| Basic | 0.0358 | 0.148 | 0.069 | 0.920 | 0.765 | 2339 |
| Logit +0.5 | 0.0212 | 0.144 | 0.064 | 0.918 | 0.751 | 2489 |
| Logit +1 | 0.0359 | 0.188 | 0.115 | 0.907 | 0.726 | 2790 |
| Logit +2 | 0.0502 | 0.243 | 0.173 | 0.847 | 0.646 | 3801 |
| Sex () | 0.0270 | 0.146 | 0.068 | 0.916 | 0.756 | 2430 |
| Sex () | 0.0305 | 0.156 | 0.079 | 0.906 | 0.743 | 2552 |
| Sex () | 0.0284 | 0.158 | 0.082 | 0.903 | 0.738 | 2605 |
| Age () | 0.0263 | 0.167 | 0.077 | 0.917 | 0.750 | 2380 |
| Age () | 0.0308 | 0.170 | 0.079 | 0.916 | 0.748 | 2403 |
| Age () | 0.0314 | 0.174 | 0.096 | 0.901 | 0.735 | 2431 |
Table 5c. (c) Trained on CheXpert and tested on MIMIC ().
Details of the Generative Model
For completeness, we summarise the semi-supervised deep structural causal model of [9], which we use throughout this work to generate counterfactual images. We refer the reader to the original paper for full derivations; here we restate only what is needed to make our experiments reproducible.
Setup
The endogenous variables comprise an image and a set of causal variables , denoted when observed and when predicted. Latent variables form part of the exogenous noise for and are modelled with a hierarchical latent structure [4, 24]. In addition to this generative backbone, the model contains a predictive component that infers from , which is what permits counterfactual generation when labels are missing. As in [9], we present the case of a single cause variable and a single effect variable ; the construction extends to any finite number of variables under an arbitrary causal structure.
Training Objective
Training proceeds by partitioning the data according to which labels are available. For fully labelled samples drawn from , the ELBO gives
where is the ELBO of a conditional VAE and are Gaussian priors. For unlabelled samples from , the labels are predicted via and regularised against their priors:
When only the cause is observed, for ,
and when only the effect is observed, for ,
Since the parent predictors and would otherwise only receive gradient when their targets are unobserved, an explicit classification term is added on the labelled subset [15], yielding the total objective
In the final term the labelled variables are placed in a topological ordering [6] beginning at the root nodes, so that for every the ancestors of lie in and its descendants in .
Two key implementation details of [9] are retained. First, for discrete variables whose true label is unavailable, the corresponding term is inversely weighted by the entropy of the predicted label distribution — e.g. the expectation in (16) is scaled by — using predictive uncertainty as a proxy for how far the imputed label should be trusted. Second, the expectations over labels in (15) are computed by summing over all values of a discrete , [15], which becomes expensive as the number of variables or classes grows. Rather than resorting to Monte-Carlo estimation, training begins on the labelled subset alone, so that the predictors are already accurate before being used to impute labels for the remainder of the data.
Counterfactual Regularisation
The model additionally employs a causal form of consistency regularisation [27], in which perturbations of the input are restricted to interventions on the DAG over the causal variables. Given an image with variables , intervening on the effect via produces with variables . If the DAG holds, the cause is invariant to this intervention, so predicting and minimising — for a suitable distance — penalises violations of that invariance; if is unobserved it is first imputed as . Intervening on the cause instead should change the effect: is compared against the counterfactual value of under the intervention on , with imputed when unlabelled. This regularises both the generative component and the causal mechanisms .
Counterfactual Generation
At inference time, counterfactuals are produced using the causal mechanisms . Abduction requires the structural assignments to be invertible in the exogenous noise, so each mechanism is encoded as an invertible map , , parameterised by the causal parents [20]. Each counterfactual value is then obtained in topological order as , and any missing is imputed beforehand with the trained predictor , as illustrated in Figure 4 (right).

Figure 4. Outline of the model of [9], reproduced here for reference. Green: observed, grey: latent, red: predicted, blue: causal generative, yellow: decoding. (left) Training; the predictions are used for decoding unless the labels are observed. (right) Counterfactual generation.
Example Counterfactuals
Below we provide some example counterfactuals produced by the model and used for testing with our proposed metrics.

Figure 5. (a) From left to right: (1) original: white, diseased, 52-year-old male (2) do(age=22) (3) do(healthy) (4) do(black) (5) do(female) (6) do(all).
(b) (1) original: black, healthy, 77-year-old female (2) do(age=20) (3) do(diseased) (4) do(white) (5) do(male) (6) do(all).
(c) (1) original: white, diseased, 52-year-old male (2) do(age=82) (3) do(healthy) (4) do(asian) (5) do(female) (6) do(all).
(d) (1) original: black, healthy, 68-year-old female (2) do(age=35) (3) do(diseased) (4) do(white) (5) do(male) (6) do(all).
Fig. 5: MIMIC-CXR Counterfactuals from model trained on 20% labels.
References
- [1]Murat Seckin Ayhan and Philipp Berens. Test-time data augmentation for estimation of heteroscedastic aleatoric uncertainty in deep neural networks. In Medical Imaging with Deep Learning, 2018.
- [2]Imon Banerjee and Et al. Ai recognition of patient race in medical imaging: a modelling study. The Lancet. Digital health, 4:e406 – e414, 2021.
- [3]Robin Blythe, Rex Parsons, Marcus E. H. Ong, and Adrian Barnett. Continuous predicted risks should be retained when deploying clinical prediction models. Journal of Clinical Epidemiology, 188:112009, 2025.DOI
- [4]Fabio De Sousa Ribeiro, Tian Xia, Miguel Monteiro, Nick Pawlowski, and Ben Glocker. High fidelity image counterfactuals with probabilistic causal models. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 7390–7425, 07 2023.arxiv.org/abs/2306.15764
- [5]Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021.arxiv.org/abs/2010.01412
- [6]Mukul Gagrani, Corrado Rainone, Yang Yang, Harris Teague, Wonseok Jeon, Herke Van Hoof, Weiliang Will Zeng, Piero Zappi, Christopher Lott, and Roberto Bondesan. Neural topological ordering for computation graphs, 2022.arxiv.org/abs/2207.05899
- [7]Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In ICML, 2017.arxiv.org/abs/1706.04599
- [8]Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1–42, 1997.pubmed.ncbi.nlm.nih.gov/9117894
- [9]Yasin Ibrahim, Hermione Warr, and Konstantinos Kamnitsas. Semi-supervised learning for deep causal generative models. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 294–303. Springer, 2024.arxiv.org/abs/2403.18717
- [10]Jeremy Irvin and Et al. Chexpert: a large chest radiograph dataset with uncertainty labels and expert comparison. AAAI’19/IAAI’19/EAAI’19. AAAI Press, 2019.arxiv.org/abs/1901.07031
- [11]Alistair E. W. Johnson, Tom J. Pollard, Seth J. Berkowitz, Nathaniel R. Greenbaum, Matthew P. Lungren, Chih ying Deng, Roger G. Mark, and Steven Horng. Mimic-cxr, a de-identified publicly available database of chest radiographs with free-text reports. Scientific Data, 6, 2019.DOI
- [12]Charles Jones, Daniel C. Castro, Fabio De Sousa Ribeiro, Ozan Oktay, Melissa McCradden, and Ben Glocker. No fair lunch: A causal perspective on dataset bias in machine learning for medical imaging, 2023.arxiv.org/abs/2307.16526
- [13]Konstantinos Kamnitsas and et. al. Ensembles of multiple models and architectures for robust brain tumour segmentation. In Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries, volume 10670 of Lecture Notes in Computer Science, pages 450–462. Springer, 2018.arxiv.org/abs/1711.01468
- [14]Masanari Kimura. Understanding test-time augmentation. arXiv preprint arXiv:2402.06892, 2024.arxiv.org/abs/2402.06892
- [15]Diederik P. Kingma, Danilo J. Rezende, Shakir Mohamed, and Max Welling. Semi-supervised learning with deep generative models.
- [16]Matt J. Kusner, Joshua Loftus, Chris Russell, and Ricardo Silva. Counterfactual fairness. In Advances in Neural Information Processing Systems, volume 30, pages 4066–4076, 2017.arxiv.org/abs/1703.06856
- [17]Yan Li, Matthew Sperrin, Darren M. Ashcroft, and Tjeerd Pieter van Staa. Consistency of ranking was evaluated as new measure for prediction model stability: longitudinal cohort study. Journal of Clinical Epidemiology, 138:168–177, 2021.
- [18]Zeju Li, Konstantinos Kamnitsas, and Ben Glocker. Analyzing overfitting under class imbalance in neural networks for image segmentation. IEEE Transactions on Medical Imaging, 40(3):1065–1077, 2021.
- [19]Lena Maier-Hein, Annika Reinke, Patrick Godau, et al. Metrics reloaded: Recommendations for image analysis validation. Nature Methods, 21:195–212, 2024.
- [20]Nick Pawlowski, Daniel C. Castro, and Ben Glocker. Deep structural causal models for tractable counterfactual inference. In Advances in Neural Information Processing Systems, 2020.
- [21]Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of Causal Inference: Foundations and Learning Algorithms. Adaptive Computation and Machine Learning. MIT Press, 2017.
- [22]Pedro Sanchez and Sotirios A. Tsafataris. Diffusion causal models for counterfactual estimation. In Bernhard Schölkopf, Caroline Uhler, and Kun Zhang, editors, Proceedings of the First Conference on Causal Learning and Reasoning, volume 177 of Proceedings of Machine Learning Research, pages 647–668. PMLR, 11–13 Apr 2022.arxiv.org/abs/2202.10166
- [23]Bonnie Sibbald and Chris Roberts. Understanding controlled trials: Crossover trials. BMJ, 316(7146):1719, 1998.pubmed.ncbi.nlm.nih.gov/9614025
- [24]Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. Florian Wenzel, Jasper Snoek, Dustin Tran, and Rodolphe Jenatton. Hyperparameter ensembles for robustness and uncertainty quantification. In Advances in Neural Information Processing Systems, volume 33, 2020.DOI
- [26]Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. In Advances in Neural Information Processing Systems, volume 33, 2020.arxiv.org/abs/2002.08791
- [27]Qizhe Xie, Zihang Dai, Eduard H. Hovy, Minh-Thang Luong, and Quoc V. Le. Unsupervised data augmentation for consistency training. arXiv: Learning, 2019.arxiv.org/abs/1904.12848
- [28]Zikang Xu, Jun Li, Qingsong Yao, Han Li, Mingyue Zhao, and S. Kevin Zhou. Addressing fairness issues in deep learning-based medical image analysis: A systematic review. npj Digital Medicine, 7(286), 2024.