Medical AI · Machine learning · Lesson 2 of 9
Building a cohort and avoiding data leakage
Reviewed by Dr. Ram Kumar G, MBBS, MD (Internal Medicine), MPH (USA), on · Last updated
- Intermediate
- Technical focus
- About 35 minutes
- Includes exercise
Before this lesson: Prediction or causation? Framing the clinical question, Healthcare data types: records, images, signals, notes and omics
In short
Data leakage happens when a model, during development, gets information it would not have at the moment of real use: a test result recorded after the diagnosis, the same patient in both training and test sets, or preprocessing done on the whole dataset. It makes models look far better than they are. Prevent it by defining a clear prediction time, using only data available before it, splitting by patient (or by time or site), and fitting every preprocessing step inside the training data only.
You will learn to
- Define a cohort with eligibility criteria, a prediction time (time zero) and an outcome window
- Recognise the common forms of data leakage in clinical datasets
- Show in code how preprocessing and row-level splits inflate performance
- Use pipelines and grouped splits to evaluate models honestly
Start with the moment of prediction
Every prediction model is used at a specific moment: at admission, at discharge, at a clinic visit. Defining that moment, called the index time or time zero, is the single most important design decision, because it decides which data are legitimately available.
| Element | What to define | Example: 30-day readmission at discharge |
|---|---|---|
| Eligibility | Who is in the cohort | Adults discharged alive from a medical ward |
| Time zero | When the prediction is made | The moment of discharge |
| Look-back window | Which past data may be used | Data recorded up to and including discharge |
| Outcome window | What counts as the outcome, and when | Unplanned readmission within 30 days after discharge |
| Unit of analysis | One row per what | One row per discharge, grouped by patient for splitting |
The common forms of leakage
Kapoor and Narayanan found data leakage behind errors in hundreds of published machine-learning studies across many fields. In clinical data, the usual forms are these:
| Type | What happens | Clinical example |
|---|---|---|
| Future information | Features recorded after time zero | Using the discharge diagnosis to predict deterioration at admission |
| Proxy for the outcome | A feature that exists because the outcome happened or was suspected | Antibiotic orders or a blood culture used to 'predict' sepsis |
| Same patient in train and test | The model recognises patients rather than learning risk | Several admissions of one patient split across both sets |
| Preprocessing on all data | Scaling, imputation or feature selection fitted before splitting | Selecting 'significant' lab tests using the whole dataset |
| Duplicates and near-duplicates | The same record, image or scan appears twice | Several slices of one CT scan in both training and test sets |
Leakage from preprocessing, in code
This example gives a model 5,000 features of pure noise for 200 synthetic patients. Nothing in the data predicts the outcome, so an honest AUROC should be about 0.5.
import numpy as np
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score
from sklearn.pipeline import make_pipeline
rng = np.random.default_rng(1)
X = rng.normal(size=(200, 5000)) # 200 patients, 5,000 features of pure noise
y = rng.binomial(1, 0.5, size=200) # outcome unrelated to any feature
# Wrong: choose the 20 "best" features using ALL the data, then cross-validate
X_selected = SelectKBest(f_classif, k=20).fit_transform(X, y)
leaky = cross_val_score(LogisticRegression(max_iter=1000), X_selected, y,
cv=5, scoring="roc_auc")
# Right: feature selection happens inside each training fold
pipeline = make_pipeline(SelectKBest(f_classif, k=20), LogisticRegression(max_iter=1000))
honest = cross_val_score(pipeline, X, y, cv=5, scoring="roc_auc")
print(f"Selection before splitting: AUROC {leaky.mean():.2f}")
print(f"Selection inside each fold: AUROC {honest.mean():.2f}")
# Selection before splitting: AUROC 0.88
# Selection inside each fold: AUROC 0.43- Choosing features with the whole dataset let the test folds influence the model, producing an AUROC of 0.88 from pure noise.
- With a pipeline, the honest estimate is close to chance (0.43 here; small samples vary around 0.5). Put scaling, imputation and selection in a pipeline so they are refitted on each training fold.
Leakage from splitting by row, in code
Many clinical datasets have several rows per patient: admissions, visits, images. If rows from the same patient fall in both training and test sets, a flexible model can recognise the patient instead of learning anything general.
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import GroupShuffleSplit, train_test_split
rng = np.random.default_rng(2)
rows = []
for patient in range(400):
traits = rng.normal(size=5) # stable patient characteristics
risk = rng.normal() # patient-level risk, NOT explained by traits
for _ in range(rng.integers(3, 8)): # several admissions per patient
features = traits + rng.normal(scale=0.1, size=5)
readmitted = rng.binomial(1, 1 / (1 + np.exp(-2 * risk)))
rows.append([patient, *features, readmitted])
df = pd.DataFrame(rows, columns=["patient", "f1", "f2", "f3", "f4", "f5", "readmitted"])
X, y, groups = df[["f1", "f2", "f3", "f4", "f5"]], df.readmitted, df.patient
def test_auroc(train_idx, test_idx):
model = RandomForestClassifier(n_estimators=200, random_state=0)
model.fit(X.iloc[train_idx], y.iloc[train_idx])
return roc_auc_score(y.iloc[test_idx], model.predict_proba(X.iloc[test_idx])[:, 1])
row_train, row_test = train_test_split(np.arange(len(df)), test_size=0.3, random_state=0)
grp_train, grp_test = next(GroupShuffleSplit(test_size=0.3, random_state=0).split(X, y, groups))
print(f"Random split by admission: AUROC {test_auroc(row_train, row_test):.2f}")
print(f"Split by patient: AUROC {test_auroc(grp_train, grp_test):.2f}")
# Random split by admission: AUROC 0.76
# Split by patient: AUROC 0.52- Splitting 1,976 admissions at random gives an AUROC of 0.76; splitting the 400 patients gives 0.52, the honest answer, because the features carry no real information about risk.
- Use GroupKFold or GroupShuffleSplit with the patient identifier whenever patients can appear more than once.
Splitting by time and by site
A random split, even by patient, assumes the future looks like the past. A temporal split, training on earlier years and testing on later ones, shows how a model copes with changes in practice, coding and patient mix. Testing on a different hospital, external validation, is stronger still. Report which kind of split each result comes from.
| Split | What it tests | When to use |
|---|---|---|
| Random by patient | Performance on similar, unseen patients | Always the minimum |
| Temporal | Robustness to change over time | Before deploying a model that will run for years |
| By site (external) | Transport to a new setting | Before claiming a model works beyond the development site |
A leakage checklist
- Is every feature available, in this form, at time zero in the real workflow?
- Could any feature exist only because the outcome happened or was suspected?
- Are all rows for a patient (and all images from a study) in the same split?
- Is every preprocessing step fitted inside the training data only, using a pipeline?
- Was the test set left untouched until the final evaluation, and used once?
- Does performance look too good to be true? Investigate before celebrating.
Report it
- TRIPOD+AI asks authors to report how the data were split, how predictors were defined in relation to the time of prediction, and how preprocessing was handled. Following it makes leakage easier to spot.
Exercise
Build a 30-day readmission cohort from Synthea data with an explicit time zero, then measure how much each kind of leakage inflates performance.
Data: Synthea synthetic encounters and observations (CSV export) (Synthea is Apache 2.0; generated records are synthetic)
- Select inpatient encounters, set time zero at discharge and label unplanned readmissions within 30 days.
- Build features using only data recorded before discharge, then add one deliberately leaky feature recorded after discharge.
- Compare test AUROC with and without the leaky feature, and with row-level versus patient-level splits.
- Move imputation and scaling into a scikit-learn Pipeline and confirm the estimate does not change between runs.
- Run a temporal split (train on earlier years, test on the latest year) and compare it with the random patient split.
The runnable notebook and worked solution will be published soon. You can work through the steps with the code on this page in the meantime.
Key takeaways
- Define time zero first; only data available before it may be used.
- Split by patient, and by time or site when you can.
- Fit all preprocessing inside the training data, using pipelines.
- Results that look too good usually are: check for leakage.
Terms used: Cross-validation, Data leakage, Dataset shift, External validation
Related: Prediction or causation?, A logistic-regression risk score, Healthcare data types, Evaluation, safety & deployment track
FAQ
What is data leakage in machine learning?
Information getting into model development that would not be available when the model is used for real, such as future data, outcome proxies or test-set information. It makes performance look better than it will be.
Why should I split by patient rather than by row?
If the same patient's records are in both training and test sets, the model can recognise the patient instead of learning general patterns, which inflates test performance.
Is it leakage to standardise features before splitting?
Yes, a mild form: the scaling uses information from the test set. Usually the effect is small for scaling, but large for feature selection or imputation, so always fit preprocessing inside a pipeline.
What is time zero?
The moment at which the prediction is made in practice, such as discharge or a clinic visit. It defines which data may be used as predictors and when the outcome window starts.
Is cross-validation enough to show a model works?
No. Cross-validation estimates performance on similar data from the same source. Temporal and external validation are needed to show the model will work over time and in other settings.
Sources
- Kapoor S, Narayanan A. Leakage and the reproducibility crisis in machine-learning-based science. Patterns 2023;4(9):100804. doi:10.1016/j.patter.2023.100804
- Kaufman S, Rosset S, Perlich C, Stitelman O. Leakage in data mining: formulation, detection, and avoidance. ACM Trans Knowl Discov Data 2012;6(4):15. doi:10.1145/2382577.2382579
- Roberts M, et al. Common pitfalls and recommendations for using machine learning to detect and prognosticate for COVID-19 using chest radiographs and CT scans. Nature Machine Intelligence 2021;3:199–217. doi:10.1038/s42256-021-00307-0
- Collins GS, et al. TRIPOD+AI statement: updated guidance for reporting clinical prediction models that use regression or machine learning methods. BMJ 2024;385:e078378. doi:10.1136/bmj-2023-078378
- scikit-learn documentation. Common pitfalls and recommended practices (data leakage)
Educational content only. This material teaches methods; it is not clinical advice. Code, models and examples here are not validated medical devices and must not be used to make decisions about real patients without proper clinical validation, governance and any required regulatory approval. Examples use synthetic or openly licensed de-identified data only. Never paste identifiable patient information into public AI tools. Mentions of models, tools or vendors are for identification, not endorsement.