sEHR-CE: Language modelling of structured EHR data for efficient and generalizable patient cohort expansion
Abstract
Electronic health records (EHR) offer unprecedented opportunities for in-depth clinical phenotyping and prediction of clinical outcomes.
Combining multiple data sources is crucial to generate a complete picture of disease prevalence, incidence and trajectories.
The standard approach to combining clinical data involves collating clinical terms across different terminology systems using curated maps, which are often inaccurate and/or incomplete.
Here, we propose sEHR-CE, a novel framework based on transformers to enable integrated phenotyping and analyses of heterogeneous clinical datasets without relying on these mappings.
We unify clinical terminologies using textual descriptors of concepts, and represent individuals’ EHR as sections of text.
We then fine-tune pre-trained language models to predict disease phenotypes more accurately than non-text and single terminology approaches.
We validate our approach using primary and secondary care data from the UK Biobank, a large-scale research study.
Finally, we illustrate in a typediabetes use case how sEHR-CE identifies individuals without diagnosis that share clinical characteristics with patients.
Introduction
Electronic health records (EHR) are collected as part of routine medical care, and include demographic information, disease diagnoses, laboratory results, medication prescriptions, etc., providing a patient’s clinical state over time.
Recent machine learning techniques have been used to exploit the richness of EHR data at scale for diagnosis, prognosis, treatment and understanding of disease [Li et al., 2020, Steinberg et al., 2021].
Many medical terminologies are used across clinical data sets and the standard practice involves mapping clinical terms across different resources [Li et al., 2020, Hassaine et al., 2020] or onto common data models [Stang et al., 2010].
Here, we propose sEHR-CE (language modelling of S tructured EHR data for patient C ohort E xpansion), a novel framework based on transformers to enable the integrated analysis of multiple clinical resources without relying on any manual curation and mapping.
Using text descriptions of concepts as input, our method generalises across data modalities and terminologies, i.e. text and structured EHR.
This enables us to leverage a plethora of pre-trained language models like PubMedBERT [Gu et al., 2022] to encode each patient’s medical record.
We ask the model to learn representations of clinical histories from diagnosed patients (cases) to predict phenotypes (such as diseases).
In the absence of a directly comparable model, we evaluate the performance of our model against that of Li et al.
[2020], as the state-of-the-art approach for learning clinical term embeddings for future disease prediction.
See Appendix A for a more complete survey.
In our validation experiments, we will show that the model can then identify individuals with missing diagnoses (controls) that share a similar clinical history with cases, indicating they might have the disease or be at risk of developing it.
The contributions of our paper are (i) the presentation of a cohort expansion method that provides phenotype predictions outperforming non-text and single terminology approaches, and (ii) an in-depth qualitative evaluation demonstrating that positively predicted controls share similar clinical representations as cases, providing a high degree of evidence that these may be previously un- or misdiagnosed individuals.
Finally, we demonstrate our method’s potential for patient stratification by disease severity.
Methods
Input Generation .
For each EHR source and associated ontology a ∈ A , we denote the set of concepts (e.g. clinical terms) as Θ a , and the set of text descriptions as Ξ a .
The total vocabulary of concepts and text descriptions across all ontologies is denoted by Θ =
(cid:83) a ∈A Θ a and Ξ =
(cid:83) a ∈A Ξ a , respectively.
For each patient, we define their full clinical history through time and across sources by the concatenation of sequences of clinical descriptions ( ξ θ 1 , . . .
, ξ θ t ) , ξ θ i ∈ Ξ , i = 1 , . . .
, t , ordered in time.
More details and an example can be found in Appendix B.1, and Figure B.1.
To form the input, we process the raw text sequence of descriptions into tokens (e.g. words and sub-words)
X = W ( ξ θ 1 , . . .
, ξ θ t ) =
( x 1 , . . .
, x n ) under a fixed size vocabulary V with the WordPiece tokenizer
W [Wu et al., 2016].
Label Generation .
Let ∆ =
( d 1 , . . .
, d D ) denote an ordered set of clinical outcomes or events, in our case
disease phenotypes d i .
For each phenotype d i ∈ ∆ , we let 1 d i be an indicator function that assigns a binary label to individuals according to the presence or absence of d i .
To define 1 d i , we rely on external oracles or phenotype definitions, such as CALIBER (Appendix B.2).
Figure 1 shows a schematic of input and label generation for a given patient.
Model Design .
Let X ( p ) =
( x ( p ) 1 , . . .
, x ( p ) n ) denote the tokenized input sequence of individual p .
It forms the input to an encoding function x ( p ) 1 , . . .
, x ( p ) n = Encoder ( X ( p ) ) , where each x i is a fixed-length vector representation of each input token x i .
Let y ( p ) =
( y ( p ) 1 , . . .
, y ( p ) D ) , y ( p ) i ∈ { 0 , 1 } , be labels denoting presence or absence of phenotypes d 1 , . . .
, d D .
Given a learned representation over inputs, we decode over y ( p ) under the predictive model P ( y ( p ) | X ( p ) ) .
We calculate the probability of each phenotype d i given the input sequence encoding P ( y ( p ) i | x ( p ) 1 , . . .
, x ( p ) n ) via a decoder module.
Specifically, we decode the representation into logits per phenotype z ( p ) 1 , . . .
, z ( p ) D =
Decoder ( x ( p ) 1 , . . .
, x ( p ) n ) and calculate the probability per phenotype as P ( y ( p ) d | x ( p ) 1 , . . .
, x ( p ) n ) =
σ ( z ( p ) d ) , where σ denotes the sigmoid function (Figure B.2).
We will omit the superscript ( p ) denoting the sample index in the remainder of the text.
2.1 Data Augmentation
Clinical Masking .
We mask input descriptions ξ θ from term-description pairs ( θ, ξ θ ) with 1 d ( θ, ξ θ ) = 1 for d ∈ ∆ using the following clinical masking strategy during training and validation as in Devlin et al.
[2019], Wei and Zou [2019]: Remove ξ θ with 80% probability, retain ξ θ with 10% probability, and replace ξ θ with a randomly selected description from the corpus with 10% probability.
Figure 2 displays a worked example.
During testing, these term-description pairs are fully removed from the input sequence.
Comorbidities .
This masking approach is straightforward if an individual has only one positive label, but many people have comorbidities, e.g. co-occurring conditions.
To allow for a sample with multiple positive labels d i 1 , . . .
, d i n , we create n input samples by replicating both the input sequence and target vector of phenotype labels.
Consider the j th replicated input sequence, and let d = d i j .
We mask this replicated input sequence with the masking strategy for phenotype d described in the previous section; e.g. ξ θ is masked if 1 d ( θ, ξ θ ) = 1 (Figure 3). 2.2 Defining a Loss Function in the Presence of Comorbidities and varying Prevalences
Loss weights .
The data augmentation approach for comorbidities may lead to our model overfitting when an input sequence contains the descriptions associated with an existing phenotype
d that is not masked.
To account for the contribution of the target vector y to the loss function in these scenario, we define a binary masking vector γ j = ( γ j 1 , . .
. , γ jD ) where γ jj = 1 and γ jk = 0 , k = 1 , . . .
, D , k ̸
= j .
Individuals with no positive phenotype labels are assigned an all-zero masking vector (Figure 3).
Then, for a given input text sequence X , target label vector y =
( y 1 , . .
y D ) and masking vector γ
= ( γ 1 , . .
. , γ D )
, we define the loss weights by setting ω d =
if y d = 1 and γ d = 0 , and ω d = 1 if y d = 0 or y d = 1 and γ d = 1 .
Positive weights .
Because some diseases are very rare, whereas others are very common, our prediction classes are expected to be highly imbalanced in the practical setting.
For cohort expansion, we want to increase recall while balancing a decline in precision.
For d ∈ ∆ , let TN d and TP d denote the total number of negative and positive examples of d in the training set.
We define the positive weight ρ d as ρ d = TN d /TP d , for all d ∈ ∆ .
Our loss function is defined as a mean-reduced binary cross-entropy loss function over phenotypes, where disease prevalence and comorbidities are handled with loss and positive weights:
[Equation]
where σ denotes the sigmoid function, ω
d the comorbidity-derived loss weight, ρ d the positive weight, and z ( p ) d the predicted probability for phenotype d ∈ ∆ for sample (e.g. individual) p .
Experiments and Results
Data.
The UK Biobank (UKBB)
[Sudlow et al., 2015] is a large-scale research study of around 500 k individuals.
We restrict the data set to those that have entries in both hospital and GP records, reducing our cohort to 155 k. To assess the quality of our model predictions, we choose four diseases that differ in terms of prevalence and clinical characteristics (Appendix C): Type 2 diabetes mellitus (T2DM),
Heart failure (HF), Breast Cancer, and Prostate Cancer.
We use validated phenotype definitions from
CALIBER Kuan et al.
[2019] to label patients with each of the diseases (Appendix C.1).
We test the performance of our model on its ability to diagnose known cases, compare it to other methods, and evaluate associations of clinical features with predictions on T2DM. sEHR-CE.
We use the pre-trained language model PubMedBERT [Gu et al., 2022] as the encoder of the tokenised input sequences of clinical term descriptions.
Since our input systematically differs from the general scientific text on which PubMedBERT was trained, we fine-tune on the maskedlanguage modeling (MLM) task using the full UKBB cohort.
The proposed model sEHR-CE uses the fine-tuned encoder and a fully connected linear layer as the decoder.
To train on the multi-label classification task of outcome prediction, we split our data set into five equally sampled folds f 0 , ...fcontaining unique patients, and mask the data according to our strategy (Section 2.1, Figure C.1).
We train a total of five models on three folds, holding back folds f i for validation and f ( i +1)5 for testing for model i , i = 1 , . . .
, 5 (Figure C.2).
All results presented are predictions on the independent test set.
For masking and training details, see Appendix C.2.
Baseline Models .
We compare the performance of our model sEHR-CE to BEHRT [Li et al., 2020], which takes a tokenised sequence of clinical terms, age and position embeddings as input.
Ontologies from hospital and GP records are mapped to CALIBER definitions, removing unmapped terms (more details in Appendix C.3).
We train five such BEHRT models to predict an individual developing the four phenotypes on the same five data splits as sEHR-CE.
Similarly, we train five sEHR-CE models restricted to CALIBER tokens (denoted sEHR-CE-codes) for comparison.
Figure C.3 shows the distributions of predicted probabilites for all phenotypes across all methods.
sEHR-CE shows the best performance across all four phenotypes in terms of recall at 0 .
5 and AUC on the test set (Table 1,
Figure C.4).
BEHRT performs slightly better than sEHR-CE-codes, indicating a benefit of adding visit position and age.
Performance varies across phenotypes, presumably due to different clinical characteristics making some diseases easier to predict than others. [Table] [Figure]
3.1 Clinical evaluation on Type 2 Diabetes Mellitus
T2DM lends itself as a use case to qualitatively evaluate the predictions of missed cases, as it is a well-studied, slowly developing disease with varying disease severity.
We used thresholds based on the 98 th , 90 th and 12 th percentiles of sEHR-CE’s predicted probabilities (Figure C.3) to define five different groups (Table C.1): regular cases, cases predicted with high probability, cases predicted with low probability, regular controls and controls predicted with high probability (missed cases).
Measures of Disease Severity .
Haemoglobin A1c (HbA1c) is a blood biomarker used to diagnose and monitor diabetes.
A level of 48mmol/mol or higher is considered diabetes; while a value between 42 and 48 mmol/mol is considered pre-diabetes.
The input data did not include such biomarker data, so we use it here for evaluation.
We aggregated HbA1c measurements taken in primary care (GP) of each individual by taking the 95 -th percentile value.
Figure 4a shows that cases predicted with high probability had the highest HbA1c mean levels.
Cases identified with low probability were in the prediabetic range of HbA1c levels, indicating a less severe state.
Missed cases had elevated
HbA1c levels close to the prediabetic stage, representing individuals at risk of developing T2DM.
We further investigated the association of the T2DM predicted probabilities with other measures of disease severity, expanded in Appendix C.4.
Taken together, our results demonstrate that sEHR-CE’s predicted probabilities of being diagnosed with T2DM are associated with disease severity.
Polygenic risk scores .
Genetic risk for complex diseases like T2DM arises from many genetic changes that, when taken together, can increase an individual’s risk of developing the disease, which can be defined by polygenic risk scores (PRS).
Sinnott-Armstrong et al.
[2021] developed PRS based on the UK Biobank participants for a set of diseases, including T2DM.
We computed and standardised
T2DM PRS across all individuals in our cohort [Lewis and Vassos, 2017]. Figure 4b demonstrates that higher predicted probability of T2DM was associated with a higher genetic risk.
Conclusion
We presented a data-driven method for cohort expansion based on language modelling.
Our approach fuses primary and secondary care data via text, and we propose a data augmentation approach to allow for comorbidities in a patient’s history.
Our method predicts disease phenotype labels more accurately than non-text and single terminology approaches.
We presented a high degree of evidence that our model identifies previously undiagnosed individuals that can extend the original cohort for downstream analysis.
Future work will consider methods that are less restrictive on sequence length [Kitaev et al., 2020, Beltagy et al., 2020] and allow for irregular time steps [Shukla and Marlin, 2021] and age [Kazemi et al., 2019], as well as adding more data sources (e.g. medications).
A Related Work
ML approaches have been applied to EHR data either using a cross-sectional matrix of diagnosis terms (e.g. an ICD-10 term)
[Henderson et al., 2017, Miotto et al., 2016] or sequential data as input [Choi et al., 2016a, Pham et al., 2016, Alaa and van der Schaar, 2019, Choi et al., 2016b, Denaxas et al., 2018, Cai et al., 2018, Darabi et al., 2020].
In the former, methods are blind to the order in which diagnoses occur and subsequently how a patient’s disease profile develops.
In the latter, most methods do not learn embeddings for the full sequence of diagnoses in a patient’s medical history, and instead learn embeddings per diagnosis term or, at most, a short sequence of terms whether they are utilising LSTMs [Choi et al., 2016a], RNNs [Pham et al., 2016], CNNs [Nguyen et al., 2017], or transformer models [Li et al., 2020].
To be able to deal with heterogeneous ontologies from different
EHR data sources, all of these models rely on noisy and often lossy mappings across ontologies or on phenotyping algorithms, eg. manually curated groupings of ontology terms, such as CALIBER [Kuan et al., 2019].
In contrast, our method uses the textual description of terms to learn representations across the full sequence of diagnoses in a patient’s history.
Recent work by Hur et al.
[2022] aims to unify
EHR records by learning description-based embeddings from multiple data sources.
Our work was developed in parallel independently and addresses the specific use case of cohort expansion, instead of merely providing examples of potential downstream applications to assess improvements of predictive power.
B Method details and figures B.1 Fusing ontologies via text
We consider all EHR data sources with their ontologies where each concept has a textual descriptor.
For example, the ontologies of GP and hospital records are made up of clinical terms (e.g. Read version 2/ Clinical Terms Version 3 and ICD9/ICD10 codes, respectively) and their description.
For each EHR source and associated ontology a ∈ A , we denote the set of concepts (e.g. clinical terms in the case of GP or hospital records) within this ontology as Θ a and the set of text descriptions as Ξ a .
The total vocabulary of concepts and text descriptions across all ontologies is denoted by Θ =
(cid:83) a ∈A Θ a and Ξ =
(cid:83) a ∈A Ξ a , respectively.
For each patient, we define their full clinical history through time and across sources as the sequence of time-indexed concepts as ( θ 1 , . . .
, θ t ) , θ i ∈ Θ , i = 1 , . . .
, t .
The premise of our work relies on the assumption that for every concept θ ∈ Θ , there exists a unique text description ξ θ ∈ Ξ .
For example, under the ICD-10 terminology, the alphanumeric code E11. has the associated description ‘Type 2 diabetes mellitus without complications’.
Thus the clinical history of each patient can be uniquely represented by the concatenation of sequences of clinical descriptions ( ξ θ 1 , . . .
, ξ θ t ) , ξ θ i ∈ Ξ , i = 1 , . . .
, t , ordered in time.
Figure B.1 shows an example of fusing primary care and hospitalization records.
B.2 Label Generation: Oracle Feature Tagging for Disease Phenotyping
Given a set of clinical terms Θ and text descriptions Ξ , we rely on external oracles to assign labels to a given set of target phenotypes ∆ .
We assume that for each phenotype d ∈ ∆ there exists a mapping 1 d : Θ × Ξ → { 0 , 1 } , ( θ, ξ θ ) (cid:55)→ δ d indicating whether the presence of d can be inferred from the clinical term and its description.
An aggregated phenotype label of 1 is assigned, if 1 d ( θ, ξ θ ) =for any of the term-description pairs ( θ, ξ θ ) in the input sequence, and 0 otherwise (Figure 1).
Here, ∆ is a set of disease phenotypes, that can be taken from disease-specific phenotyping algorithms [Chapman et al., 2021], such as CALIBER [Kuan et al., 2019], which consists of manually-curated sets of clinical terms across primary and secondary care ontologies for defining 308 chronic and acute disease phenotypes.
B.3 Model Design
Figure B.2 shows a diagram of sEHR-CE.
C Experiment design and model predictions
To assess the quality of our model predictions, we chose four diseases that differ in terms of prevalence and clinical characteristics.
Type 2 diabetes mellitus (T2DM) is one of the most prevalent chronic diseases worldwide [WHO], and patients are primarily diagnosed and managed in primary care (GP)
[Medscape].
Heart failure (HF) is one of the main causes of death in the older population [Ferreira et al., 2019], and its acute manifestations are treated in hospital care.
Malignant neoplasms of the breast and of the prostate are both less prevalent diseases almost exclusively present in only biologically females or males, respectively [Ly et al., 2013, Rawla, 2019].
We test the performance of our model on its ability to diagnose known cases, compare it to other methods, and evaluate associations of clinical features with the predictions on T2DM with available orthogonal data.
C.1 Data Processing
The UK Biobank (UKBB)
[Sudlow et al., 2015] is a large-scale research study of around 500 k individuals between the ages of 40 and 54 at the time of recruitment.
It includes rich genotyping and phenotyping data, both taken at recruitment and during primary and secondary care encounters (GP and hospital).
We use patient records from GP and hospital visits in the form of code ontologies
Read version2/ Clinical Terms Version 3, and ICD-9/10 together with their textual descriptions.
We restrict the data set to individuals that have entries in both hospital and GP records, which reduces our cohort to 154 , 668 individuals.
Requiring individuals to have entries in their GP records reduces bias towards acute events that usually present in hospitals, but we note that removing individuals without any hospital records may still bias the data towards more severe cases.
We use CALIBER, previously validated phenotype definitions from Kuan et al.
[2019] to label patients with T2DM, HF, malignant neoplasm of the breast, and malignant neoplasm of the prostate.
A patient can be admitted to the hospital for multiple days.
We treat an entire hospital admission as one point in time using the admission date, and only keep unique ICD-10/ICD-9 codes for each visit.
We aggregate visits that are less than a week apart into one visit keeping only unique codes.
This approach removes repeated codes, thus avoiding redundancy and reducing sequence length.
Only patients with at least 5 clinical terms present in their clinical history are included to allow for sufficient information for any predictions, reducing the cohort to a final 129 , 932 individuals.
We use phenotype definitions from CALIBER
[Kuan et al., 2019] to label patients with T2DM, HF, malignant neoplasm of the breast, and malignant neoplasm of the prostate.
Each phenotype definition consists of a list of ICD10 and Read2, Read3 codes and their children.
The usage of PubMedBERT restricts the length of input sequences we can use.
To avoid excluding relevant clinical information by truncating the input sequences, we break up patient histories longer than the limit into multiple input sequences of smaller length with the same target vector.
C.2 Model Training
We use the pre-trained language model PubMedBERT [Gu et al., 2022] as the encoder of the tokenised input sequences of clinical term descriptions.
Since our input systematically differs from the general scientific text on which PubMedBERT was trained, we fine-tuned on the masked-language modeling (MLM) task (Figure B.2, Figure C.1), by masking words (e.g. descriptions) at random following the original BERT paper [Devlin et al., 2019].
The model, fine-tuned using the full UKBB cohort of 138 , 079 patients, was trained with early stopping for 5 epochs with a batch size of 32 and a learning rate of 4 × 10 − 5 using gradient descent with an AdamW optimizer, and weight decay of 0 .
The output dimension of the encoder was 768 .
The proposed model sEHR-CE is using the fine-tuned encoder and a fully connected linear layer as the decoder.
To train on the multi-label classification task of outcome prediction, we split our data set into five equally sampled folds f 0 , ...
f 4 containing unique patients, using a stratified sampling method to maintain the same phenotype proportion in every split [Sechidis et al., 2011], and mask the data according to our strategy (Section 2.1, Figure C.1).
We train a total of five models for 3 epochs on three folds, holding back folds f i for validation and f ( i +1)5 for testing for model i , i = 1 , . . .
, 5 (Figure C.2).
This maintains a 60 / 20 / 20 training, validation and testing split overall while providing us with enough training and testing examples.
All results presented are predictions of each model on its respective independent test set.
We used a learning rate of 10 − 5 , and a warm-up proportion of 0 . 25 .
Performance was monitored every 0 .epochs on the validation fold.
C.3 Model evaluation
As explained in 3, we compare the performance of our model sEHR-CE to BEHRT
[Li et al., 2020].
BEHRT takes a tokenised sequence of clinical terms, age and position embeddings as input.
Ontologies from hospital and GP records are mapped to CALIBER definitions [Kuan et al., 2019], removing unmapped terms.
Phenotype definitions in CALIBER include different categories (for example, ‘diabetes’ contains categories ‘type 1’ and ‘type 2’), that were ignored by the original
BEHRT publication, so we expanded the token set to define a token per CALIBER phenotype and category.
A transformer model is pre-trained to predict masked tokens before it is trained to predict a set of possible diagnoses an individual may develop given the input sequence.
Similarly, we trained five sEHR-CE models restricted to CALIBER tokens (denoted sEHR-CE-codes) for comparison.
All results presented are predictions of each model on its respective independent test set.
Figure C.3 shows the predicted probabilities for cases and controls across all phenotypes and models, and figure C.4 shows the AUPRC curves for each phenotype and method in the test sets.
To avoid inclusion of too many false positives, we defined missing cases as those controls with predicted probability in the 98 th percentile.
C.4 Expanded Qualitative Evaluation of T2DM Diagnosis Prediction
Five groups were used to evaluate the model’s T2DM diagnosis predictions.
Table C.1 shows the percentiles of sEHR-CE’s predicted probabilities to define each group, along with each size.
[Table]
[Table]
[Table]
We then investigated the association of predicted probability and several proxies of disease severity: number of GP and hospital admissions , survival and risk of cardiovascular disease.
C.4.1 Number of GP visits and hospital admissions
As expected, both cases and controls with a high predicted probability of a T2DM diagnosis, exhibit a slightly higher number of GP and hospital visits than the other groups (Figure C.5), indicating that they are experiencing a more severe form of T2DM requiring care.
This is particularly higher in the case of hospital visits, indicating patients experiencing acute events: both cases and controls with a high predicted probability visit a hospital approximately 10 times more often than their low probability counterparts.
Although the model was not given information from which data source the input data was coming from, this analysis indicates that it has learned to associate acute events with disease severity.
C.4.2 Survival analysis
To compare survival across different groups of individuals, we use the Kaplan-Meier estimator with all-cause mortality as the endpoint with right-censored data accounting for individuals without any event occurrence since the last follow-up.
Both cases and controls with high predicted probabilities had the lowest survival, followed by general cases, controls and finally cases with low predicted probability (Figure C.6a), indicating that the model’s predicted probability is associated with survival.
C.4.3 Cardiovascular Risk
T2DM is a known risk factor and comorbidity of cardiovascular disease, which, in turn, is the most prevalent cause of death in T2DM patients.
The GP records contain Framingham and QRISK scores; these are two scores that assess an individual’s risk of developing cardiovascular disease [Figure] within the next 10 years, based on several coronary risk factors.
The Framingham score is derived from an individual’s age, gender, total cholesterol, high-density lipoprotein cholesterol, smoking habits, and systolic blood pressure, whereas the QRISK score, which is almost exclusively used now, extends this score with additional factors such as body mass index, ethnicity, measures of deprivation, chronic kidney disease, rheumatoid arthritis, atrial fibrillation, diabetes mellitus, and anti-hypertensive treatment.
Both cases and controls with high predicted probability of having T2DM had a higher risk of developing cardiovascular disease compared to their low predicted probability counterparts (Figure C.6b) indicating that the model has learned to associate the risk of developing both diseases at the same time.