Medical AI · Machine learning · Lesson 6 of 9
Missing data and class imbalance
Reviewed by Dr. Ram Kumar G, MBBS, MD (Internal Medicine), MPH (USA), on · Last updated
- Intermediate
- Technical focus
- About 30 minutes
- Includes exercise
Before this lesson: A logistic-regression risk score, Healthcare data types: records, images, signals, notes and omics
In short
In clinical data, whether a value is missing often depends on how ill the patient seemed, so dropping incomplete records or ignoring missingness can bias a model. For prediction, imputation inside the modelling pipeline, sometimes with a 'was it measured?' flag, usually works best, as long as testing practices stay stable. Rare outcomes rarely need 'correcting': oversampling or reweighting leaves AUROC unchanged and makes predicted risks far too high.
You will learn to
- Distinguish data missing completely at random, at random, and not at random
- Compare dropping, imputing and flagging missing values in a prediction model
- Explain when missingness indicators help and when they make a model fragile
- Show why class-imbalance corrections damage calibration, and what to do instead
Why values are missing
Rubin's classification describes three mechanisms. The mechanism matters more than the amount of missing data.
| Mechanism | Meaning | Clinical example |
|---|---|---|
| Missing completely at random (MCAR) | Missingness unrelated to anything | A sample lost because a courier was late |
| Missing at random (MAR) | Missingness explained by recorded data | HbA1c missing more often in younger patients, and age is recorded |
| Missing not at random (MNAR) | Missingness depends on the unrecorded value or on unrecorded reasons | Lactate measured mainly when a clinician suspects sepsis |
Informative missingness, in code
In this simulation, lactate is mostly measured in sicker patients, as often happens in practice.
import numpy as np
import pandas as pd
rng = np.random.default_rng(13)
n = 8000
severity = rng.normal(0, 1, n)
lactate = (1.5 + 0.8 * severity + rng.normal(0, 0.5, n)).clip(0.3).round(1)
died = rng.binomial(1, 1 / (1 + np.exp(-(-3.0 + 1.6 * severity))))
measured = rng.binomial(1, 1 / (1 + np.exp(-(-0.5 + 1.8 * severity)))).astype(bool)
df = pd.DataFrame({"age": rng.integers(18, 90, n),
"heart_rate": (85 + 6 * severity + rng.normal(0, 12, n)).round(),
"lactate": np.where(measured, lactate, np.nan),
"died": died})
print(f"Lactate missing in {df.lactate.isna().mean():.0%} of patients")
print(df.groupby(df.lactate.isna()).died.mean().rename("death rate").round(3))
# Lactate missing in 58% of patients
# lactate
# False 0.184
# True 0.042
# Name: death rate, dtype: float64- Patients without a lactate result died at a quarter of the rate of those with one (4.2% versus 18.4%). Whether the test was done carries information.
- Training only on complete records would keep 42% of patients, all from a much sicker group, and produce a model that overestimates risk for everyone else.
Dropping, imputing and flagging
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
X, y = df[["age", "heart_rate", "lactate"]], df.died
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3,
random_state=0, stratify=y)
def evaluate(columns, imputer, label):
model = make_pipeline(imputer, StandardScaler(), LogisticRegression(max_iter=1000))
model.fit(X_train[columns], y_train)
p = model.predict_proba(X_test[columns])[:, 1]
print(f"{label:<34} AUROC {roc_auc_score(y_test, p):.3f}")
evaluate(["age", "heart_rate"], SimpleImputer(), "Lactate dropped")
evaluate(["age", "heart_rate", "lactate"], SimpleImputer(strategy="median"),
"Median imputation")
evaluate(["age", "heart_rate", "lactate"],
SimpleImputer(strategy="median", add_indicator=True),
"Median imputation + missing flag")
# Lactate dropped AUROC 0.658
# Median imputation AUROC 0.768
# Median imputation + missing flag AUROC 0.784- Imputation inside the pipeline keeps every patient and uses the lactate values that exist. Adding a missingness flag lets the model use the fact that the test was or was not ordered.
- The imputer is fitted on the training data only, as part of the pipeline, so no test information leaks in.
A flag makes the model depend on clinicians' habits
- The missingness flag works because clinicians order lactate when worried. If testing practice changes, for example a new sepsis protocol that measures lactate in everyone, the flag's meaning changes and the model can fail silently.
- Use flags knowingly, document them, and monitor testing rates after deployment.
Imputation for prediction is not the same as for research
For estimating associations or treatment effects, multiple imputation with the outcome included in the imputation model is the standard approach under the missing-at-random assumption. For prediction, the imputation must be something that can be applied at the moment of use, when the outcome is unknown. Sperrin and colleagues argue that these two goals need different handling.
| Goal | Recommended approach | Why |
|---|---|---|
| Describe or estimate effects | Multiple imputation, including the outcome; pool estimates | Gives unbiased estimates and honest uncertainty under MAR |
| Predict for new patients | Deterministic imputation fitted on training data (median, regression or iterative), plus flags if missingness will be stable | Must be reproducible at prediction time without the outcome |
| Tree ensembles | Native missing-value handling | The model learns where to send missing values |
Class imbalance: usually not a problem to fix
Many clinical outcomes are rare, and it is common to 'balance' the data by oversampling cases, undersampling non-cases, generating synthetic cases (SMOTE) or reweighting. For risk prediction, these corrections usually do harm: they leave discrimination unchanged and make the model overestimate risk.
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split
rng = np.random.default_rng(4)
n = 20000
X = rng.normal(size=(n, 4))
y = rng.binomial(1, 1 / (1 + np.exp(-(-3.5 + X @ np.array([0.9, 0.6, 0.4, 0.0])))))
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3,
random_state=0, stratify=y)
print(f"Outcome rate {y.mean():.1%}")
for label, weights in [("No correction", None), ("class_weight='balanced'", "balanced")]:
model = LogisticRegression(class_weight=weights).fit(X_train, y_train)
p = model.predict_proba(X_test)[:, 1]
print(f"{label:<24} AUROC {roc_auc_score(y_test, p):.3f} "
f"mean predicted risk {p.mean():.1%} observed {y_test.mean():.1%}")
# Outcome rate 4.7%
# No correction AUROC 0.811 mean predicted risk 4.7% observed 4.7%
# class_weight='balanced' AUROC 0.811 mean predicted risk 37.4% observed 4.7%- Reweighting did not improve AUROC at all, but the average predicted risk jumped from 4.7% to 37.4% for an outcome that occurs in 4.7% of patients.
- If the concern is that a 0.5 threshold flags almost no one, the solution is to choose a threshold that suits the clinical decision, not to distort the probabilities.
- Van den Goorbergh and colleagues showed this in simulations and real data: imbalance corrections led to strongly miscalibrated models without better discrimination.
Checklist
- Report how much data is missing for each variable, and think about why.
- Avoid complete-case analysis unless missingness is small and plausibly random.
- Put imputation inside the modelling pipeline; plan how values will be imputed at prediction time.
- Use missingness flags only if testing practices will stay stable, and monitor them.
- Do not rebalance classes for risk prediction; choose decision thresholds instead, and report calibration.
Exercise
Handle informative missingness and a rare outcome properly, and measure the consequences of common shortcuts.
Data: Synthetic patients generated in the notebook (no real patient data) (Generated in the notebook)
- Reproduce the lactate example and add a complete-case model; compare its mean predicted risk with the observed rate in the full test set.
- Replace median imputation with IterativeImputer and compare AUROC and calibration.
- Simulate a change in practice in which lactate is measured in everyone, and evaluate the model with the missingness flag on this new data.
- Reproduce the imbalance example, then try random oversampling and SMOTE (imbalanced-learn); compare AUROC, calibration plots and Brier scores.
- For the uncorrected model, choose a threshold that gives 80% sensitivity and report the resulting PPV.
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
- Missingness in clinical data is often informative; think about why values are missing.
- For prediction, impute inside the pipeline in a way that can be repeated at prediction time.
- Missingness flags can help but tie the model to current testing habits.
- Do not rebalance rare outcomes for risk models; choose thresholds and check calibration.
Terms used: Calibration, Class imbalance, Data leakage, Missing data mechanisms (MCAR, MAR, MNAR)
Related: Healthcare data types, Diagnostic accuracy metrics, Gradient-boosted trees for readmission risk
FAQ
Should I drop patients with missing data?
Usually not. Complete-case analysis discards information and, when missingness depends on illness severity, trains the model on an unrepresentative group.
Is mean or median imputation acceptable?
For prediction, simple imputation inside a pipeline can work reasonably well, especially combined with a missingness flag. For estimating associations or effects, multiple imputation is preferred.
What is a missingness indicator?
An extra variable that records whether a value was missing. It lets the model use the information carried by whether a test was done, at the cost of depending on testing habits.
Should I use SMOTE for rare outcomes?
For clinical risk prediction, generally no. SMOTE and similar methods tend to make models overestimate risk without improving discrimination. Choose a decision threshold instead.
How do I pick a threshold for a rare outcome?
Base it on the clinical decision: the balance of harms from false positives and false negatives, the capacity to act on alerts, and the resulting sensitivity and PPV. Decision-curve analysis can help.
Sources
- Rubin DB. Inference and missing data. Biometrika 1976;63(3):581–592. doi:10.1093/biomet/63.3.581
- Sterne JAC, et al. Multiple imputation for missing data in epidemiological and clinical research: potential and pitfalls. BMJ 2009;338:b2393. doi:10.1136/bmj.b2393
- Sperrin M, et al. Missing data should be handled differently for prediction than for description or causal explanation. J Clin Epidemiol 2020;125:183–187. doi:10.1016/j.jclinepi.2020.03.028
- van den Goorbergh R, et al. The harm of class imbalance corrections for risk prediction models: illustration and simulation using logistic regression. J Am Med Inform Assoc 2022;29(9):1525–1534. doi:10.1093/jamia/ocac093
- Agniel D, Kohane IS, Weber GM. Biases in electronic health record data due to processes within the healthcare system. BMJ 2018;361:k1479. doi:10.1136/bmj.k1479
- scikit-learn documentation. Imputation of missing values
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.