Making Conformal Predictors Robust in Healthcare Settings: a Case Study on EEG Classification

AI summary

Overview

Research area: Uncertainty quantification for clinical machine learning, specifically conformal prediction applied to electroencephalogram (EEG) classification.

Technical level: Intermediate. The paper assumes familiarity with conformal prediction, calibration/test splits, and coverage guarantees, but the case study structure keeps the empirical narrative accessible.

Scope: An empirical and partially theoretical evaluation of four conformal prediction methods on two EEG benchmarks from the TUH EEG Corpus, showing that personalized calibration (NCP) improves coverage substantially under random splits but that all methods fail to reach target coverage under patient-level distribution shift.

What This Paper Is About

Standard conformal prediction produces prediction sets with theoretical coverage guarantees, but those guarantees assume the calibration and test data are independent and identically distributed. In healthcare, and especially in EEG, that assumption breaks down because patients differ from one another and labels themselves are uncertain (expert neurologists often disagree, and labels come from group voting). This paper asks whether conformal prediction methods can be made more robust to patient-level distribution shift, and it tests that question on EEG seizure and event classification.

Key Contributions

1. A neighborhood conformal prediction (NCP) approach that improves coverage by up to 20% over baselines, as stated in the paper's contribution list.
2. Theoretical results showing that NCP yields provably better coverage under covariate shift, with further theoretical results deferred to reference [1].
3. Evidence that prediction set sizes remain relatively stable under personalization, rather than growing to compensate for improved coverage.
4. A modular integration into PyHealth, an open-source healthcare AI framework, enabling replication in other clinical settings.

Main Findings

• Covariate CP does not improve coverage over Naive CP. The two methods track similarly across all α values and both datasets, because likelihood-ratio estimation via kernel density estimation (KDE) is unreliable in high-dimensional EEG feature spaces.

• NCP substantially improves coverage on the random split. On TUEV under the random split, NCP achieves around 34% greater coverage than Naive CP at α = 0.2 (0.87 vs. 0.53), with a modest increase in average prediction set size (1.22 vs. 0.90).

• TUAB is the easier case. On TUAB, Naive CP already satisfies coverage because the task is binary, so the personalization benefit is less pronounced, though NCP still produces notably more compact prediction sets.

• Patient-level splits expose the limits of all approaches. Making calibration and test patients disjoint causes every method to fall well short of target coverage on TUEV, with high variance across seeds, and a smaller but still present gap on TUAB.

• NCP's tight sets may be a liability under strong shift. NCP's smaller set sizes under the patient split suggest its personalization yields overly tight prediction sets when distribution shift is severe.

• Empirical coverage remains imperfect. No method consistently achieves 1 − α coverage across all settings, and cross-patient distribution shift remains an open challenge for conformal prediction in EEG classification.

• Abstract-level claim vs. headline result. The abstract states that personalized calibration strategies can improve coverage by over 20 percentage points while maintaining comparable prediction set sizes; the introduction states NCP improves coverage by up to 20% over baselines; the TUEV α = 0.2 comparison reports 0.87 vs. 0.53.

Methodology in Plain English

The researchers took two benchmark tasks derived from the TUH EEG Corpus: TUEV, a six-class EEG event classification task covering SPSW, GPED, PLED, EYEM, ARTF, and BCKG; and TUAB, a binary normal-versus-abnormal detection task. They trained ContraWR, a ResNet-based 2D CNN operating over multi-channel spectrograms, from scratch across five random seeds.

They then compared four conformal prediction methods. Naive CP is standard split conformal prediction assuming i.i.d. calibration and test data. Covariate CP reweights using a likelihood ratio estimated with KDE. K-means CP derives its threshold from the nearest cluster's calibration samples. NCP performs sample-specific calibration via k-nearest neighbors with relevance weighting. Naive CP and Covariate CP are non-personalized; K-means CP and NCP are personalized.

The key experimental design choice is the splitting regime. Under the random split, all samples are pooled and divided globally at 60%, 10%, 15%, and 15% for training, validation, calibration, and test. Under the patient split, patients are partitioned first, the train partition is further split at 60%, 20%, and 20% for train, validation, and calibration, and held-out patients form the test set — which introduces realistic cross-patient distribution shift. Comparing the two regimes isolates how much of the failure comes from patient-level shift rather than from the conformal method itself.

Why This Matters

Impact on research. The paper provides a controlled, seed-replicated case study showing that personalized conformal methods help under mild (random) shift but do not rescue coverage under patient-level shift. It also shows that some theoretically motivated corrections, like covariate reweighting via KDE, fail in practice at the dimensionality of EEG features. The negative result is as informative as the positive one, and the code being available through PyHealth lowers the barrier to reproducing and extending it.

Real-world applications:

• Clinical decision support for EEG review, where a prediction set can flag which waveforms a model is uncertain about rather than forcing a single label.
• Seizure and abnormality triage, where calibrated uncertainty could route ambiguous recordings to human neurologists instead of automated pipelines.
• Deployment monitoring, where a sudden drop in empirical coverage could serve as an early warning that the incoming patient population has shifted away from the calibration population.
• Generalization to other clinical modalities — ECG, imaging, or electronic health records — that face the same patient-level shift problem.

Industry relevance. Any regulated or safety-critical healthcare AI product needs defensible uncertainty estimates, not just accuracy numbers. This paper shows that the choice between non-personalized and personalized conformal calibration changes coverage by large margins on a real clinical task, which directly affects how much a deployed system can be trusted. The PyHealth integration (pip install pyhealth) means these methods can be dropped into existing pipelines rather than reimplemented.

Future Directions

1. Applying personalized conformal predictors across a broader range of healthcare tasks to identify which types of distribution shift these methods handle well and which they do not.
2. Improving NCP's coverage by using stronger EEG foundation model embeddings, such as TFM-Tokenizer, which better characterize signal structure.
3. Addressing the specific failure mode where NCP's personalization produces overly tight prediction sets under strong cross-patient shift, since tight sets that miss the true label are worse than uninformative ones.
4. Closing the empirical coverage gap under patient-level splits, which no method tested here achieved on TUEV, and which showed high variance across seeds.

Target Audience

Researchers and practitioners working on uncertainty quantification for clinical machine learning, especially those applying conformal prediction to physiological time-series data such as EEG. It is also relevant to clinical AI engineers evaluating whether coverage guarantees will hold once a model leaves a random-split evaluation and meets real patients, and to readers interested in a concrete example of where a theoretically sound correction (KDE-based covariate reweighting) breaks down in high-dimensional feature spaces.

Authors’ abstract

Quantifying uncertainty in clinical predictions is critical for high-stakes diagnosis tasks. Conformal prediction offers a principled approach by providing prediction sets with theoretical coverage guarantees. However, in practice, patient distribution shifts violate the i.i.d. assumptions underlying standard conformal methods, leading to poor coverage in healthcare settings. In this work, we evaluate several conformal prediction approaches on EEG seizure classification, a task with known distribution shift challenges and label uncertainty. We demonstrate that personalized calibration strategies can improve coverage by over 20 percentage points while maintaining comparable prediction set sizes. Our implementation is available via PyHealth, an open-source healthcare AI framework: https://github.com/sunlabuiuc/PyHealth.