Source-linked AI summary
Counterfactual Marginalisation: Framework for Evaluating Robustness to Nuisance Variables
Yasin Ibrahim, Hermione Warr, Robin J. Evans, Konstantinos Kamnitsas
TL;DR
Medical classifiers can achieve strong held-out performance while relying on demographic or acquisition-related shortcuts, motivating evaluation beyond standard point estimates. The paper uses counterfactual image generation and intervention-aware prediction averaging to define robustness metrics, finding that these metrics expose demographic sensitivity and predictor bias better than standard evaluation.
Problem
Strong held-out performance does not guarantee clinically appropriate decision rules because models may rely on demographic, acquisition-related, or site-specific shortcuts.
Method
The framework intervenes on nuisance parent variables in SCM-generated counterfactual images, averages predictions over a target intervention distribution, and defines metrics for risk, calibration, stability, and worst-case sensitivity.
Results
Counterfactual metrics expose predictor biases better than standard risk estimation and measure nuisance-variable sensitivity without requiring disease labels.
Takeaways & Limitations
Counterfactual evaluation provides intervention-aware robustness measures that can reveal demographic sensitivity when standard point-estimate metrics appear stable.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper proposes counterfactual marginalisation as a test-time evaluation procedure for detecting classifier reliance on demographic and other nuisance variables. It introduces intervention-aware metrics and shows they can expose predictor biases better than traditional metrics without requiring disease labels.
- 1 Introduction: High held-out accuracy may conceal reliance on demographic, acquisition-related, or site-specific shortcuts instead of pathology-relevant evidence.Causal modelling motivates evaluating how predictions change when demographic attributes are intervened upon while patient-specific factors remain fixed.
- 1 Introduction: Counterfactual marginalisation evaluates medical image classifiers by intervening on nuisance variables at test time.The procedure extends counterfactual use beyond explanation and augmentation to robustness evaluation.
- 1 Introduction: The framework introduces metrics for counterfactual risk, interventional expected risk, stability, and worst-case counterfactual risk.
- 1 Introduction: The approach exposes predictor biases better than traditional metrics and measures sensitivity to nuisance variables without disease labels.
2 Background
Counterfactual marginalisation is related to methods that improve robustness by averaging over uncertainty or transformed inputs. Its causal analogue averages predictions over counterfactual changes rather than parameters or generic image transformations.
- 2 Background: Counterfactual marginalisation is framed as a causal analogue of ensembles, Bayesian model averaging, and test-time augmentation.
3 Counterfactual Marginalisation for Evaluation
Counterfactual marginalisation evaluates classifiers by intervening on nuisance parent variables, generating counterfactual images, and averaging predictions over a target intervention distribution. It extends this evaluation to risk, calibration, stability, decision flips, and ranking robustness.
- Counterfactual prediction: Counterfactual images intervene on selected parent variables while preserving each sample’s inferred exogenous noise, enabling intervention-aware predictions for the same patient-specific factors.The method partitions parents into variables to marginalise, such as age and sex, and variables kept fixed, then uses abduction–action–prediction to generate counterfactuals.
- Counterfactual prediction: The target intervention distribution q(e_a) can represent the empirical test distribution, uniform demographic strata, or an external population distribution.
- Counterfactual risk: Counterfactual marginal risk evaluates performance after marginalising nuisance variables, while interventional expected risk averages loss across interventions before aggregation.
- Counterfactual risk: Worst-case sensitivity is assessed with CVaR over the highest-loss interventions, including the risk of the single worst intervention.CVaR measures expected loss within the worst α-fraction of interventions under q(e_a).
- Calibration and stability: Counterfactual calibration uses Brier score and ECE on marginalised probabilities, while stability metrics quantify output variation, threshold flips, observational–marginalised disagreements, and rank changes.The rank-stability measure compares observed and marginalised patient rankings using Spearman correlation.
4 Theoretical Advantages of Counterfactual Evaluation
Counterfactual estimation can improve efficiency for subgroup quantities by using every test patient while preserving individual-level causal information. Its variance advantage grows as the target subgroup becomes rarer.
- The CF estimator uses counterfactual samples for every test patient, rather than only members of the observed subgroup, to estimate subgroup means.For an attribute value z, it generates a counterfactual image at A = z for each patient and averages the predictions.
- The estimator's variance is bounded by σ2(1 − n_z/n), with σ2 denoting the observational subgroup variance.The result follows from conditioning on abducted exogenous noise and applying the Rao-Blackwell theorem.
- As the subgroup proportion n_z/n approaches zero, the variance bound approaches the generation-noise term, increasing the relative advantage for rare subgroups.The paper connects this efficiency gain to within-subjects designs, which reuse information from the same individuals across conditions.
5 Experiments
Experiments on CheXpert and MIMIC evaluate counterfactual metrics across demographic interventions, showing that they track predictor bias more reliably than original-data risk and support calibration and stability assessment.
- The experiments use balanced CheXpert and MIMIC-CXR test sets for pleural-effusion classification and controlled counterfactual interventions on sex and age.Interventions include two sexes and eight ages from 20 to 90 years, producing K = 16 counterfactual images per patient under a uniform distribution.
- Counterfactual metrics better discern increasing predictor bias than standard risk, which can decrease as bias increases.In CheXpert, standard cross-entropy risk sometimes falls from Logit+0.5 to Logit+1 despite greater bias.
- Marginalised ECE increases more consistently with predictor bias than original ECE, suggesting improved calibration assessment.
- Stability metrics are evaluated against worst-case AUC gaps across sex, race, and age stratifications for unbiased and biased predictors.
- Stability metrics correlate strongly with AUC gaps, suggesting they can serve as proxies when test-time labels are unavailable.
- Across the cross-dataset experiment, increasing bias generally raises S_var, S_flip, and D_obs while lowering S_rank, supporting their use for estimating bias levels.
6 Conclusion
The conclusion presents counterfactual marginalisation as a test-time framework for evaluating medical image classifiers under controlled nuisance-variable interventions. It reports that counterfactual evaluation exposes sensitivity to demographic interventions across CheXpert and MIMIC.
- Counterfactual marginalisation averages predictions over SCM-generated variants to measure risk, calibration, stability, tail sensitivity, and ranking robustness.
A Full Results
The results section compares marginal risk and calibration across CheXpert and MIMIC settings, including cross-dataset evaluation. The reported tables organize original and marginalised risk and calibration measures.
- A.1 Risk: Table 3 reports marginal risk decomposition under a uniform prior, including original, complete-marginalisation, individual-environment, tail, worst-case, and transition-count measures.The caption defines CTC as correct transition count and marks higher CTC as better.
- A.1 Risk: The risk results include CheXpert-trained and tested models, MIMIC-trained and tested models, and models trained on CheXpert but tested on MIMIC.The corresponding sample sizes are 5,180 for CheXpert and 28,189 for MIMIC.
- A.2 Calibration: Table 4 compares original and marginalised calibration using Brier score, expected calibration error, and maximum calibration error across 10 bins.The table distinguishes marginalised predictions from original predictions.
A.3 Stability
The stability section evaluates prediction changes under marginalisation using score, label-flip, distance, rank-correlation, and rank-change measures. It reports results for CheXpert, MIMIC, and cross-dataset testing.
- A.3 Stability: Table 5 evaluates stability at threshold τ = 0.5 under a uniform prior using prediction variance, label-flip rate, score difference, rank correlations, and mean absolute rank change.The listed rank measures are Spearman ρ and Kendall τb.
- A.3 Stability: The stability evaluation covers models trained and tested on CheXpert, trained and tested on MIMIC, and trained on CheXpert but tested on MIMIC.The CheXpert and MIMIC test sets contain 5,180 and 28,189 samples, respectively.
- A.3 Stability: The counterfactual-generation model represents images with observed and latent variables and includes a predictive component for inferring labels when they are missing.Its hierarchical latent structure supplies exogenous noise for the image.
B.2 Training Objective
Training partitions data by label availability and combines variational objectives, supervised classification terms, uncertainty weighting, exact discrete expectations, and causal consistency regularisation. Inference generates counterfactuals through invertible causal mechanisms and imputation of missing variables.
- B.2 Training Objective: Training uses separate objectives for fully labelled, unlabelled, and partially labelled samples, with predicted labels used when observations are missing.The partially labelled case observes only the effect variable, while the total objective adds an explicit classification term on labelled data.
- B.2 Training Objective: The objective places labelled variables in topological order so each variable’s ancestors precede it and descendants follow it.This ordering is used in the final objective term.
- B.2 Training Objective: Unobserved discrete labels are uncertainty-weighted, and expectations over discrete variables are computed by summing across all possible values.Training starts on labelled data so predictors are accurate before imputing labels for the remaining data.
- B.2 Training Objective: Causal consistency regularisation penalises violations of invariance when intervening on an effect and checks corresponding effect changes when intervening on a cause.The regulariser applies to both the generative component and causal mechanisms.
- B.2 Training Objective: At inference, counterfactuals are generated in topological order using invertible causal mechanisms, with missing variables imputed by the trained predictor beforehand.Abduction preserves the individual’s exogenous noise for counterfactual generation.
C Example Counterfactuals
The paper illustrates counterfactual images generated by intervening on age, disease status, race, sex, or all listed attributes. These examples are drawn from MIMIC-CXR counterfactuals generated with 20% labels.
- C Example Counterfactuals: The counterfactual examples were produced by the model and used for testing with the proposed metrics.A separate figure caption identifies MIMIC-CXR counterfactuals generated from a model trained on 20% labels.
- C Example Counterfactuals: The examples show original images alongside interventions on age, disease status, race, sex, and all attributes.The depicted interventions include age changes to 20, 22, or 82 years and changes in demographic or disease attributes.
- C Example Counterfactuals: The example sets include a white diseased 52-year-old male and a black healthy 77-year-old female before counterfactual interventions.The first and third example sets begin from the same white diseased 52-year-old male.