A framework for understanding label leakage in machine learning for health care

Abstract

Introduction

The pitfalls of label leakage, contamination of model input features with outcome information, are well established. Unfortunately, avoiding label leakage in clinical prediction models requires more nuance than the common advice of applying “no time machine rule.”

Framework

We provide a framework for contemplating whether and when model features pose leakage concerns by considering the cadence, perspective, and applicability of predictions. To ground these concepts, we use real-world clinical models to highlight examples of appropriate and inappropriate label leakage in practice.

Recommendations

Finally, we provide recommendations to support clinical and technical stakeholders as they evaluate the leakage tradeoffs associated with model design, development, and implementation decisions. By providing common language and dimensions to consider when designing models, we hope the clinical prediction community will be better prepared to develop statistically valid and clinically useful machine learning models.

Introduction

Numerous challenges prevent the successful development, integration, and diffusion of machine learning (ML) solutions in health care.^1,2 Label leakage—when outcome labels are subtly included in or leaked into model inputs—is one challenge that can undermine model implementations.³ For example, concern may arise when antihypertensive medications are used to predict hypertension⁴; antibiotics are used to predict sepsis⁵; or termination of mechanical ventilation or entry into hospice care are used to predict death.⁶ While these cases highlight seemingly clear leakage, the complexities of clinical care actually make some leaks appropriate, even necessary. However, label leakage critiques often fail to consider the nuances of clinical care, the diagnostic process, and how ML tools are used in practice.

In this perspective, we address this gap by presenting a novel conceptual framework for ML practitioners to better assess model designs and implementations for appropriate and inappropriate label leakage. We ground our discussion in a variety of real-world ML models and apply the framework to determine the validity of various model inputs. We discuss the implications of our framework to ensure clinical and technical stakeholders understand the tradeoffs associated with model design decisions. We conclude with recommendations for ML practitioners.

Label leakage definitions

Label leakage, also known as target or data leakage, is a broad term constituting any contamination of model input features or predictors with outcome information. There are 3 primary types of label leakage—aggregation-based, timing-based, and association-based.

Aggregation-based leakage is caused by data preprocessing and partitioning. Site-level label leakage may occur when data across multiple sites is aggregated for model training. The model may learn that certain sites are more highly associated with an outcome and correlate site characteristics with the label.⁷ Similarly, patient-level label leakage may occur when the same patient is represented in both training and validation datasets. The model may associate the label with specific patients and fail to generalize to unseen patients.

Timing-based leakage is caused by including predictor information generated after the time of prediction, thus violating the “no time machine rule.”^3,8 Avoiding these temporal errors sounds simple but is complicated by latencies that delay the availability of health data. For example, diagnosis codes for an inpatient admission are often finalized weeks after discharge. Blood culture results are often finalized days after collection. Timing-based label leakage could easily occur by using encounter diagnosis codes or blood culture results in an inpatient deterioration model. On the other hand, clinically acceptable violations of the “no time machine rule” can occur when patients and providers are interested in predicting risk under alternative potential treatment plans.⁹ The planned treatments are not recorded past events, but planned events that would occur in the future prior to outcome ascertainment.

Association-based leakage occurs when features available at the time of prediction are deterministically related to the outcome. Such proxies can act as surrogate outcomes and although these features may not violate the “no time machine rule,” their inclusion in a prediction model could undermine model utility. For example, a pneumothorax ML model was trained on a cohort of chest X-rays where a substantial portion of positive cases included chest tubes that were visible in the image.¹⁰ On a separate test set, the model performed poorly on pneumothorax cases that were not already recognized by clinicians and treated with a chest tube. Unfortunately, this type of label leakage is poorly characterized.³ Identifying association-based leakage requires close inspection of how model inputs and outputs are related, as well as nuanced understanding of a model’s intended use.

Unpacking deterministic relationships

Deterministic relationships between model inputs and outcomes are difficult to parse. Clinical diagnosis is a dynamic process of refinement from suspicion to collection of evidence through intervention to reassessment to confirmation.¹¹ However, ML tools often assume a simple, linear process: input data elements are fed into an ML algorithm to predict an outcome label which then prompts a clinician to consider an intervention. As such, ML tools largely fail in “helping clinicians navigate the dynamic refinement process of diagnosis.”¹¹ Critically, data generated during this iterative diagnostic process complicates the evaluation of label leakage and raises the concern of including model inputs that may be deterministically related to the outcome. As information accumulates through the diagnostic process, informative predictive model inputs may be collected to improve model performance and surrogate outcomes may be captured causing label leakage. Model developers must carefully account for the dynamic nature of healthcare data collection and known clinical evaluation processes as they design ML tools.

Let us consider 2 examples. First, a patient at risk of sepsis is administered antibiotics out of an abundance of caution. Is the antibiotic administration a leaky predictor for sepsis? Second, a decompensated patient in the intensive care unit dies shortly after discontinuation of mechanical ventilation. Is discontinuation of mechanical ventilation a leaky predictor for death? In both cases, one could argue to remove the potentially leaky predictors, but that approach is misguided. If antibiotic administrations or mechanical ventilation changes are unobserved, the ML model will almost certainly learn confounded relationships between predictors (eg, vital signs, analyte results) and outcomes. A more nuanced approach is needed to identify problematic label leakage while ensuring that ML solutions are clinically useful. This is because the choice to include these variables impacts the applicable use of each model, rather than generate label leakage resulting in globally unusable models.

A framework for assessing label leakage

Below, we propose a framework that considers the cadence, perspective, and applicability of predictions. Figure 1 illustrates these concepts and their relation to critical temporal windows of patient care. Placing a model within this framework provides guidance to distinguish appropriate and inappropriate label leakage and resulting limitations of clinical use. Table 1 and Figures 2 and 3 provide examples drawn from real ML projects.

Figure 1.

Framework for evaluating label leakage over key temporal periods and their interaction with model contextual factors of cadence of predictions, perspective of predictions, and facets of prediction applicability.

Table 1.

Use cases drawn from real clinical machine learning projects highlighting important design considerations when evaluating label leakage.

Model name	Outcome predicted	Eligibility criteria	Perspective of prediction	Conditions to initiate predictions	Conditions to terminate predictions	Frequency of predictions
Cadence: cross-sectional models
Inpatient mortality	In-hospital mortality	Adult presenting to an ED and admitted to the hospital	Reactive	Time of inpatient bed request order	DNR documented	Single
Surgery complications	30-day post-operative complications	Adult presenting a surgery clinic for evaluation before undergoing an invasive procedure.	Proactive (user supplied CPT code for anticipated procedure)	Time of outpatient surgery clinic visit	None	Single
Cadence: cohort models
Sepsis watch	Sepsis	Adult presenting to an ED	Reactive	Time of ED arrival	Time of sepsis, death, or discharge from ED or hospital	Hourly
ED triaging	Admission to general inpatient bed or ICU	Adult presenting to an ED	Reactive	Time of ED arrival	Time of inpatient bed request order, discharge from ED, or death	Hourly
Immune-related adverse events	ED visit; Hospital admission	Adult receiving immunotherapy treatment	Proactive (using scheduled chemotherapy treatment course)	Time of immunotherapy treatment course initiation	6 months since immunotherapy treatment course initiation; Time of death, ED visit, or hospital admission	Weekly
Kidney failure risk population rounding	2-year risk of end stage renal disease	Adult enrolled in the Duke Connected Care Medicare Shared Savings Program	Reactive	Time of program enrollment	Time of transfer out of program; Time of death	Monthly

Abbreviations: CPT, Current Procedural Terminology; ED, Emergency Department; ICU, intensive care unit.

Figure 2.

Example use cases of the label leakage evaluation framework applied to cross-sectional models. Abbreviations: ED, Emergency Department; DNR, do not resuscitate order.

Figure 3.

Example use cases of the label leakage evaluation framework applied to cohort models.

Cadence of predictions

Models generally fall into 2 categories, which we call cross-sectional and cohort models. Cross-sectional models generate a single prediction when a patient meets pre-specified conditions. Model activity ceases until the patient next meets the pre-specified conditions. For example, an inpatient mortality model generates a single prediction when adult patients are admitted from a Duke University Medical Center Emergency Department (ED).¹² Cohort models generate repeated predictions while a patient meets pre-specified conditions. For example, a sepsis model implemented at Duke University Medical Center generates predictions every hour while adult patients are in the ED.¹³

Cohort models are more likely than cross-sectional models to include inappropriate timing-based and association-based label leakage. This is because when a model is used at a single time point, assessment of leakage is straightforward. As described earlier, clinical diagnosis is an iterative process of information gathering, diagnostic assessment, and interventions.⁸ When an ML model repeatedly generates predictions, the cohort model may include intermediate interventions as inputs prior to diagnostic certainty.¹⁴ Responses to treatments and interventions themselves can up-weight or down-weight the diagnostic probability of a condition or outcome, and potentially result in changes in therapy as well.⁸ Thus, prior to developing and implementing a cohort model, it is critical to determine the exact combination of measurements and interventions that converge to provide a surrogate for the outcome. Providing explicit bounds on the usable window of each model is also critical, enabling careful assessment of label leakage and allowable features, as well as limits of implementations. As noted in Figure 1, using data recorded prior to prediction and any surrogate outcomes would prevent inappropriate label leakage.

Perspective of predictions

Models can be designed as reactive, to support interventions that respond to historical and current patient status, or proactive, to inform future care planning. Reactive models predict the probability of an outcome and must strictly follow the “no time machine” rule—only accessing information that would be available at the time of prediction. On the other hand, proactive models that support future care planning intentionally violate the “no time machine” rule—by design, accessing information about future planned interventions that would not be available in the medical record at the time of prediction. For example, surgical risk models often utilize the planned procedure Current Procedural Terminology (CPT) code as a model input to assess risk of post-operative complications for a specific planned procedure.⁹ Similarly, an immunotherapy model considers the planned chemotherapy treatment regimen as a model input to assess risk of unanticipated hospital admission during the treatment course.¹⁵ If planned interventions change, a new prediction may be generated using all data currently available.

Applicability of predictions

It is important to distinguish between 2 facets of the applicability of predictions generated by an ML model—statistical validity and clinical utility. From a statistical perspective, as long as the information presented to a model is not deterministic of the outcome (ie, a surrogate outcome) and is available at the time of prediction, then the data and predictions are valid. However, statistically valid predictions may have no clinical value. At some point during the iterative clinical diagnosis process, further predictions lose clinical relevance, even if the outcome or surrogate outcome are not met. For this reason, some data scientists recommend the cessation of predictions at a certain point as the patient approaches the outcome,⁸ despite model inputs being statistically valid. When predictions lose clinical relevance, which is dependent on use case context, then further predictions could be stopped. For example, a mortality model used by intensive care unit clinicians to escalate care loses clinical relevance the moment a patient has mechanical ventilation discontinued. However, the same model at the same moment may remain clinically useful to organ transplant surgeons.

Recommendations

Framing model development, evaluation, and implementation with the above language can help distinguish inappropriate from appropriate label leakage. Below, we offer recommendations to ML practitioners.

Model developers, evaluators, and clinician users should collaborate closely. Avoiding inappropriate label leakage requires partnerships between technical and clinical teams throughout the model development process. Interdisciplinary teams can proactively define the cadence, perspective, and applicability of ML model predictions to help avoid inappropriate label leakage. Key decisions to document include critical time windows and criteria for statistical validity and clinical utility; combinations of interventions and clinical measures that serve as surrogate outcomes; and whether or not the model needs to account for future planned interventions. To support transparency and transportability, we suggest explicitly recording leakage-avoidance design choices, defining relevant surrogate outcomes, and documenting when model outputs are not statistically valid or clinically useful in a warnings section of Model Facts labels.¹⁶

Preserve physiologic dynamics between model features. Clinical interventions with known impacts on physiologic variables included in an ML model should be included as model inputs when they would not constitute association-based label leakage as proxy outcomes. For example, consider a sepsis model that includes temperature as an input. If the intended use of the sepsis model is to prompt consideration of antibiotic administration, then antibiotics would be a surrogate outcome. In this case, antibiotics should not be included as inputs. On the other hand, if the intended use of the model is to ensure completion of sepsis bundles, then antibiotic administrations should be included as inputs. As long as the prediction remains clinically useful, inclusion of clinical interventions that avoid label leakage can enhance statistical validity by preventing spurious correlations.

Enforce carefully defined temporal windows when building development datasets. When building inputs for reactive models, only data available at the time of prediction should be curated to avoid time-based label leakage. It is particularly important to note when laboratory and procedural information becomes available in the medical record and to associate results with the appropriate timestamps. For proactive models, similar considerations apply, with the exception of planned future interventions.

Evaluate cohort models over various times throughout episodes of care. Cohort models that generate predictions throughout an episode of care should be evaluated at various times during the episode to ensure both statistically validity and clinical utility are sustained. Results of such evaluations may inform clinical use. For example, a sepsis model may perform much better during the initial portion of hospital encounters and performance may deteriorate after long inpatient stays. Such evaluations may reveal unanticipated label leakage from outcome surrogates overlooked during model planning or deemed acceptable under a different use case design than eventually implemented. For example, a model originally intended for decision support aimed at ensuring a sepsis bundle was completed could allow antibiotic administration as an input feature without label leakage. However, leakage occurs if that model is then used in a decision support tool targeting antibiotic administration and evaluations over episodes of care may highlight this misalignment.

Always apply statistically valid models, but focus implementation design on clinical utility. Implementations should prioritize clinical utility by disabling predictions when they no longer provide actionable information, even when the predictions are free from label leakage and remain statistically valid. If a patient has yet to meet the outcome or surrogate outcome definitions but all recommended interventions for the outcome have been completed, predictions are still statistically valid but no longer clinically useful. Minimize reminders triggered by model predictions in such cases.

Conclusion

Avoiding inappropriate label leakage is much more complex than simply adhering to the “no time machine rule.” A novel framework for ML models helps distinguish appropriate from inappropriate label leakage. Use of this framework can help ensure that ML models are both statistically valid and clinically useful.

Author contributions

All authors conceived of and debated the concepts, definitions, and framework described. S.E.D. and M.P.S. lead the writing effort, with contributions and insight from M.E.M. and S.B. All authors approved the submitted manuscript.

Funding

This work received no specific funding from any agency in the public, commercial, or not-for-profit sectors.

Conflicts of interest

M.P.S. and S.B. are inventors of intellectual property licensed by Duke University to Clinetic, Inc, and Cohere-Med, Inc. M.P.S. and S.B. hold equity in Clinetic, Inc. M.E.M. and S.E.D. have no conflicts of interest to disclose.

Data availability

Not applicable. No data were generated in the course of this work.