Context Clues: Evaluating Long Context Models for Clinical Prediction Tasks on EHRs

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Wornow, Michael, Bedi, Suhana, Hernandez, Miguel Angel Fuentes, Steinberg, Ethan, Fries, Jason Alan, Re, Christopher, Koyejo, Sanmi, Shah, Nigam H.
Natura: Preprint
Pubblicazione: 2024
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866910882760491008
author Wornow, Michael
Bedi, Suhana
Hernandez, Miguel Angel Fuentes
Steinberg, Ethan
Fries, Jason Alan
Re, Christopher
Koyejo, Sanmi
Shah, Nigam H.
author_facet Wornow, Michael
Bedi, Suhana
Hernandez, Miguel Angel Fuentes
Steinberg, Ethan
Fries, Jason Alan
Re, Christopher
Koyejo, Sanmi
Shah, Nigam H.
contents Foundation Models (FMs) trained on Electronic Health Records (EHRs) have achieved state-of-the-art results on numerous clinical prediction tasks. However, most existing EHR FMs have context windows of <1k tokens. This prevents them from modeling full patient EHRs which can exceed 10k's of events. Recent advancements in subquadratic long-context architectures (e.g., Mamba) offer a promising solution. However, their application to EHR data has not been well-studied. We address this gap by presenting the first systematic evaluation of the effect of context length on modeling EHR data. We find that longer context models improve predictive performance -- our Mamba-based model surpasses the prior state-of-the-art on 9/14 tasks on the EHRSHOT prediction benchmark. For clinical applications, however, model performance alone is insufficient -- robustness to the unique properties of EHR is crucial. Thus, we also evaluate models across three previously underexplored properties of EHR data: (1) the prevalence of "copy-forwarded" diagnoses which creates artificial repetition of tokens within EHR sequences; (2) the irregular time intervals between EHR events which can lead to a wide range of timespans within a context window; and (3) the natural increase in disease complexity over time which makes later tokens in the EHR harder to predict than earlier ones. Stratifying our EHRSHOT results, we find that higher levels of each property correlate negatively with model performance, but that longer context models are more robust to more extreme levels of these properties. Our work highlights the potential for using long-context architectures to model EHR data, and offers a case study for identifying new challenges in modeling sequential data motivated by domains outside of natural language. We release our models and code at: https://github.com/som-shahlab/long_context_clues
format Preprint
id arxiv_https___arxiv_org_abs_2412_16178
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Context Clues: Evaluating Long Context Models for Clinical Prediction Tasks on EHRs
Wornow, Michael
Bedi, Suhana
Hernandez, Miguel Angel Fuentes
Steinberg, Ethan
Fries, Jason Alan
Re, Christopher
Koyejo, Sanmi
Shah, Nigam H.
Machine Learning
Artificial Intelligence
Computational Engineering, Finance, and Science
Foundation Models (FMs) trained on Electronic Health Records (EHRs) have achieved state-of-the-art results on numerous clinical prediction tasks. However, most existing EHR FMs have context windows of <1k tokens. This prevents them from modeling full patient EHRs which can exceed 10k's of events. Recent advancements in subquadratic long-context architectures (e.g., Mamba) offer a promising solution. However, their application to EHR data has not been well-studied. We address this gap by presenting the first systematic evaluation of the effect of context length on modeling EHR data. We find that longer context models improve predictive performance -- our Mamba-based model surpasses the prior state-of-the-art on 9/14 tasks on the EHRSHOT prediction benchmark. For clinical applications, however, model performance alone is insufficient -- robustness to the unique properties of EHR is crucial. Thus, we also evaluate models across three previously underexplored properties of EHR data: (1) the prevalence of "copy-forwarded" diagnoses which creates artificial repetition of tokens within EHR sequences; (2) the irregular time intervals between EHR events which can lead to a wide range of timespans within a context window; and (3) the natural increase in disease complexity over time which makes later tokens in the EHR harder to predict than earlier ones. Stratifying our EHRSHOT results, we find that higher levels of each property correlate negatively with model performance, but that longer context models are more robust to more extreme levels of these properties. Our work highlights the potential for using long-context architectures to model EHR data, and offers a case study for identifying new challenges in modeling sequential data motivated by domains outside of natural language. We release our models and code at: https://github.com/som-shahlab/long_context_clues
title Context Clues: Evaluating Long Context Models for Clinical Prediction Tasks on EHRs
topic Machine Learning
Artificial Intelligence
Computational Engineering, Finance, and Science
url https://arxiv.org/abs/2412.16178