[Paper Review] Feature Robustness in Non-stationary Health Records: Caveats to Deployable Model Performance in Common Clinical Machine Learning Tasks

Bret Nestor, Matthew B. A. McDermott+6 more|arXiv (Cornell University)|Aug 2, 2019

Machine Learning in HealthcareComputer Science42 references54 citations

TL;DR

The paper shows that ML models trained on historical EHR data struggle to generalise to future data due to non-stationarity, but aggregating features into clinical concepts notably reduces performance decay.

ABSTRACT

When training clinical prediction models from electronic health records (EHRs), a key concern should be a model's ability to sustain performance over time when deployed, even as care practices, database systems, and population demographics evolve. Due to de-identification requirements, however, current experimental practices for public EHR benchmarks (such as the MIMIC-III critical care dataset) are time agnostic, assigning care records to train or test sets without regard for the actual dates of care. As a result, current benchmarks cannot assess how well models trained on one year generalise to another. In this work, we obtain a Limited Data Use Agreement to access year of care for each record in MIMIC and show that all tested state-of-the-art models decay in prediction quality when trained on historical data and tested on future data, particularly in response to a system-wide record-keeping change in 2008 (0.29 drop in AUROC for mortality prediction, 0.10 drop in AUROC for length-of-stay prediction with a random forest classifier). We further develop a simple yet effective mitigation strategy: by aggregating raw features into expert-defined clinical concepts, we see only a 0.06 drop in AUROC for mortality prediction and a 0.03 drop in AUROC for length-of-stay prediction. We demonstrate that this aggregation strategy outperforms other automatic feature preprocessing techniques aimed at increasing robustness to data drift. We release our aggregated representations and code to encourage more deployable clinical prediction models.

Motivation & Objective

• Motivate and quantify how temporal drift affects clinical prediction models trained on de-identified EHR data.
• Evaluate robustness of multiple input representations, models, and training regimes to year-to-year shifts.
• Identify mitigation strategies that improve generalisation across future data batches.

Proposed method

• Use year-of-care data from MIMIC-III under a Limited Data Use Agreement to evaluate temporal generalisation.
• Compare four representations: Raw, PCA, CUI Code Spanning, and Clinical Aggregations.
• Benchmark four models: LR, RF, LSTM, GRU-D, under three training regimes (Year-Agnostic, Prior Year, Full History).
• Assess two binary tasks: ICU mortality and long length-of-stay, on the first 24 hours of ICU data.
• Measure performance with AUROC and report year-to-year drops in a non-stationary setting.

Experimental results

Research questions

• RQ1 How does model performance degrade when trained on historical data and evaluated on future data in MIMIC-III?
• RQ2 Which input representations and training regimes best withstand clinical concept drift?
• RQ3 Do clinically-motivated feature aggregations improve robustness compared to automatic preprocessing methods?
• RQ4 What is the impact of EHR system changes (CareVue to MetaVision) on predictive performance across years?

Key findings

• Raw representations exhibit substantial performance drops after the 2008 EHR transition (up to 0.29 AUROC loss for mortality with RF).
• Clinically aggregated representations reduce AUROC drop to 0.06 (mortality) and 0.03 (LOS) under full-history training.
• PCA and automatic CUIs provide limited robustness to year-to-year shifts, with clinical aggregation outperforming them across models.
• GRU-D shows some recovery over time, suggesting missing-data handling aids drift robustness, but aggregation remains superior.
• Year-agnostic evaluation overestimates performance, especially for raw features, highlighting deployability gaps for public benchmarks.
• Model performance saturates quickly for these tasks, suggesting a shift to harder problems or new signals beyond 181 lab/vital features.