Improving the Performance of Models Trained on Small EHR-Derived Samples by Leveraging External Data with Continual Learning Methods

Abstract

The performance of an EHR-based deep learning model trained on a small sample can be improved if more data is collected. Instead of collecting more data, the model can be trained on additional data from an analogous external source. However, this risks the model learning patterns in the external data that do not generalize to the target sample. Furthermore, data use agreements often prohibit combining datasets with medical records of different sources. We consider utilizing pre-existing methods in continual learning, namely the elastic weight consolidation (EWC) loss function and variational continual learning (VCL), both of which are regularization-based methods that we use to borrow external data and incorporate parameters from a model on external data into local model training. To investigate the utility of this modeling framework, we consider two binary classification tasks: (1) predicting which children will be diagnosed with autism spectrum disorder (ASD) from medical claims up to 18 months, and (2) predicting which patients with end-stage renal disease (ESRD) will be re-hospitalized within 30 days. Target datasets were derived from Duke University’s EHR warehouse, and external datasets were sourced from either NC Medicaid claims for the ASD prediction task, or the United States Renal Data System (USRDS) for the rehospitalization prediction task. For both of these tasks, borrowing models - using either the EWC loss function or VCL - performed similarly to that of a model trained only on the full external data, when the sample size of target data used to train the model was small. That is, while a model that does not borrow using our methods performed poorly in low data regimes, the borrowing model instead matched the performance of a model trained on external data even when sample size of target data was small. In addition, an analysis of model predictions showed that models with small samples are better calibrated and more functionally similar to a model trained only on external data when the sample size is small.

Introduction

Deep learning models trained on health system electronic health records (EHR) data have demonstrated exceptional performance across a range of clinical predictive tasks^1,2. These models perform better when they are trained on more data³. However, for many tasks, the size of the training sample may be limited, either due to the nature of the target clinical population or the data available from the healthcare system itself. In these settings, collecting additional data may not be feasible⁴.

When one does not have enough training source data, incorporating data from an external population can supplement the limited target data — provided that the datasets share comparable features and outcome definitions. These external datasets, such as those from other institutions or broader population-level studies, can provide valuable additional information for model training. However, naively introducing external data to the model may lead to poor generalization if differences in the marginal distribution of features or the conditional distribution of features given the outcome, between the external and target populations^5,6.

To carefully control the influence of a model that combines external data during the training process on a target of interest, we considered utilizing methods developed in the context of continual learning, which addresses the challenges of balancing historical data with newly observed data⁷. These methods are designed to balance the retention of learned knowledge from one dataset while effectively integrating new information from another.

In this paper, we compare two approaches for combining data sources: elastic weight consolidation (EWC)⁸ and variational continual learning (VCL)⁹. Continual learning methods, such as these, introduce regularization that encourages the model to retain knowledge from previous tasks while learning new ones. This is done to avoid what is referred to as “catastrophic forgetting.” We seek to extend this to EHR-based predictive modeling on small target samples, utilizing these regularization-based method to incorporate parameters from a model on a large external data source. In particular, we compare this approach to naïve models as the size of the training sample varies to assess whether this method provides added value. We apply these methods across two different predictive modeling tasks: predicting a future autism spectrum disorder (ASD) diagnosis, and predicting hospital readmissions among patients with end-stage renal disease (ESRD).

Methods

Data Sources

For our two tasks, the source data consist of EHR data from the Duke University Health System (DUHS). DUHS is a medium-sized healthcare system in central North Carolina. Of note, it serves as the primary provider for individuals living in Durham, NC, an until recently majority-minority county¹⁰. EHR data are stored in a research-ready datamart, mapped to the PCORnet common data model¹¹. For the external data, we used administrative data available from public payers. For the ASD prediction task, we used North Carolina Medicaid claims data. For the ESRD task, we use Medicare claims data organized by the United States Renal Data System (USRDS)¹².

Target Cohort Creation for ASD Prediction

Our ASD prediction task aims to identify children likely to receive a future diagnosis of ASD. Operationally, the model will serve as a screening tool at a child’s 18-month well child visit. For our DUHS cohort, we identified all children, born between January 1, 2014 and December 31, 2023 who had at least one well-child visit between 12 and 24 months. We used the well-child visit closest to 18-months as the anchor encounter. All International Classification of Diseases (ICD-9 and ICD-10) codes available prior to this encounter served as the predictor variables. ICD codes were mapped to clinical classification software (CCS) categories, generating 300 unique categories¹³. For the outcome, data were available through December 31, 2023. To be considered as having ASD, a child had to receive one of the ICD codes listed in Table 2 in the appendix, at two separate encounters. The date of the first code served as the child’s date of ASD diagnosis. Children without a diagnosis were censored at their last known encounter date.

Table 2.

Demographic Information for Data Used in the ESRD 30-Day Rehospitalization Prediction Task.

Data Source	Duke EHR Data	USRDS Database	SMDs
Sample Size	18236	603384
# of CCS Codes	454129	9067907
Diagnoses (Codes Per Person)	23 [18, 29]	15 [12, 18]	1.2902
Outcome Rate	4572 (25.07%)	191816 (31.79%)	0.149
Sex			0.073
Female	7930 (43.49%)	284410 (47.14%)
Race-Ethnicity
NH White	4762 (26.11%)	264795 (43.88%)
NH Black	12289 (67.38%)	232604 (38.55%)
Hispanic	645 (3.53%)	81385 (13.49%)
Asian/PI	156 (0.86%)	16139 (2.67%)
Other	384 (2.105%)	8461(1.40%)
Age	61 [51, 70]	62 [52, 71]	0.099

External Cohort Creation for ASD Prediction

For the claims-based cohort which we used as an external data source for borrowing, data were available from January 1, 2014, and December 31, 2022. We identified all children who were eligible for NC Medicaid within 30-days of birth and had continuous coverage through 18 months, with the 18-month date serving as the anchor point. We extracted diagnostic information available prior to this date. To diagnose ASD, children similarly had to have an appropriate ICD code across two encounters. Non-diagnosed children were censored either at the end of the database or when they lost NC Medicaid coverage.

Target Cohort Creation for ESRD Prediction

Our ESRD hospitalization prediction task aims to identify adults previously hospitalized with ESRD likely to be re-hospitalized after 30 days of being discharged. We identified individuals age 18 or older with ESRD, who we defined as having the ICD code N18.6X upon encounter and CPT codes 90935 or 90937 for hemodialysis, and were hospitalized at any DUHS hospital with a discharge date between January 1, 2016 to December 1, 2018. Individuals were required to be discharged to their home. We extracted all ICD codes to serve as predictors, as well as demographic factors from the index encounter. ICD codes were mapped to the CCS categories as above. The outcome was a new admission within 30 days of the discharge.

External Cohort Creation for ESRD Prediction

For the USRDS based cohort which we used as an external data source for borrowing, we identified all hospitalization encounters and individuals discharged to their home between January 1, 2016 and December 1, 2018. Individuals were required to be at least 18 years old, and to have Medicare as the primary payer at time of discharge.

Model Formulation

We considered utilizing EWC and VCL, two methods developed in the context of continual learning. Each of these methods “borrow” information from an external source to improve predictive performance on “target data”. The EWC model was initially developed to handle catastrophic forgetting, a phenomenon in which a model trained on multiple tasks rapidly underperforms in prior tasks⁹. The model adjusts a neural network’s loss function by adding a sum of squared term between the network parameters between models of a current task and a previous task, scaled by the both the prior model’s Fisher information matrix diagonals and a tunable constant hyperparameter term influencing the degree of importance of the previous task. The idea behind this modification to the network’s loss function is to induce regularization between the parameters of both models. When building a network model for clinical prediction, we can borrow external data by modeling on the target data using the EWC loss function while considering a model on external data as a previous task. In this way, borrowing from a single external source can be framed a 2-stage continual learning problem, which is non-traditional adaptation of the model rarely discussed in informatics literature. For such a model, the EWC loss function is defined as:

$$L_{\mathrm{E}\mathrm{W}\mathrm{C}}(\theta )=L_{\text{Target}}(\theta )+\frac{\lambda }{2}\sum _{i=1}^n{F}_{ii,\text{External}}{\left({\theta }_{i,\text{Target}}-{\theta }_{i,\text{External}}\right)}^2$$ (1)

Where L is a loss function, the subscript T denotes the target dataset, and the subscript E denotes the external dataset. Here, the EWC loss function adjusts the loss function of the model on target data by adding an additional set of operations, where λ is a tunable hyperparameter determining the importance of the external model fit, F_{ii, External} are the diagonals of the Fisher information matrix on model parameters from the model on External data, only, θ_{i, Target} are the model parameters of a network model trained on target data, and θ_{i, External} are the model parameters of a network model trained on external data. Note that to enable this calculation, the number of nodes and layers in the model on external data must be identical to that of the model trained on target data.

VCL is another method to facilitate information borrowing by conducting approximate Bayesian inference for neural network models⁸. It finds the posterior distribution of a network’s parameter as it learns new tasks. Since exact Bayesian inference may be intractable in most cases, VCL approximates the posterior distributions of parameters through minimizing Kullback-Leibler (KL) divergence over parameter distributions¹⁴. When building a Bayesian network model for clinical prediction, VCL can be used to first learn a model on external data, and then use the posterior distribution of this model’s parameters as the priors for a model on target data. In this way, the VCL procedure can be thought of as a way to find the common parameter space between two models. As with our non-traditional approach using the EWC loss function where we frame borrowing external data as a 2-stage continual learning problem, the VCL objective is defined as:

$$L\left(q_{\text{Target}}\left(\theta \right)\right)=\sum _{i=1}^{N_{\text{Target}}}{E}_{\theta \sim q_{\text{Target}}\left(\theta \right)}\left(\mathrm{l}\mathrm{n}p\left(y_{\text{Target}}\mid \theta ,X_{\text{Target}}\right)-KL\left(q_{\text{Target}}\left(\theta \right)\Vert q_{\text{External}}\left(\theta \right)\right)\right.$$ (2)

Where q_Target (θ) is the approximated posterior distribution of the model on target data, q_External(θ) is the approximated posterior distribution of the model on external data, N_Target is the number of observations in the model on target data, $KL\left({\mathrm{q}}_{\text{Target}}\left(\theta \right)\Vert {\mathrm{q}}_{\text{External}}\left(\theta \right)\right)$ is the Kullback-Leiber divergence measuring the difference between the posteriors of the target and external models, and y_Target and X_Target are the outcomes and features of the i^th datapoint in the model on target data.

Model Building and Assessment

When we utilize EWC and VCL to create borrowing models, we learn model parameters that are shared between the target data and external data. For this reason, we expect the models to behave as either a model trained on target data only, a model trained on external only, or somewhere in between. For each of the two predictive modeling tasks, ASD prediction and ESRD rehospitalization prediction, we construct four types of multilayer perception (MLP) models¹⁵. Each model was tested on the same held-out test data set from the target. The models were: (1) An MLP model trained only on the external data; (2) MLP models trained only on the target data; (3) MLP models that combined external and target data via EWC loss; (4) MLP models that combined external and target data via VCL. To further understand the effect of target sample size for models (2) – (4), we lowered the size of the EHR training set in intervals of 500 samples for the ASD task and in intervals of 200 samples for the ESRD task by randomly sampling-out observations, fitting a new model at each step.

The single data models were trained using a binary cross entropy loss and an Adam optimizer. The EWC borrowing models were trained using an EWC loss instead, and the VCL borrowing models were trained through variational inference. The learning rate, activation function, the number of hidden layers, and the number of nodes in each hidden layer were all fine-tuned using 10-fold grid search cross validation. For the EWC borrowing models, the λ parameter was also included in the 10-fold grid search cross validation procedure. MLP models were built using Pytorch 2.1¹⁶. Bayesian MLP models were built using Pyro 1.8.9¹⁷.

For each model, we calculated the Area Under the Receiver Operator Characteristic (AUROC) and calibration slope^{18, 19}. Furthermore, we conducted 1000 bootstrap replicates to calculate 95% confidence intervals for the AUROC estimates²⁰. Finally, to further understand the mechanism behind data borrowing we assessed the similarities of the different model functions. Specifically, we assessed the correlation between the predictions generated by all four model types, across different sample sizes. We hypothesized that under smaller sample sizes, the functional forms of the borrowing models and the generated predictions would be most similar to the external model, while under larger sample sizes the functional forms would be most similar to the models trained only on target data.

Results

Cohort Creation

We identified four cohorts, one of the target data and one of the analogous data from an external source, for each of the two tasks. The demographic characteristics of the four cohorts are shown in Table 1, along with standardized mean differences (SMDs). For the ASD task, outcome rate and sex distribution were largely similar between both target and external datasets, though there were some small differences in the distribution of race-ethnicity. For the ESRD task, the outcome rate is much higher than that of the ASD task. For this task, the outcome rate, sex distribution, and race-ethnicity distribution between both target and external datasets contained some noticeable differences.

Table 1.

Demographic Information for Data Used in the Childhood ASD Prediction Tasks

Data Source	Duke EHR Data	NC Medicaid Claims	SMDs
Sample Size	45266	426525
# of CCS Codes	1674649	26686886
Diagnoses (Codes per Person)	31 [22, 44]	39 [27, 58]	0.244
Outcome Rate	951 (2.10%)	8535 (2.00%)	0.007
Sex			0.003
Female	22138 (48.91%)	209368 (49.07%)
Race-Ethnicity			0.270
NH White	17395 (38.43%)	147487 (34.57%)
NH Black	12551 (27.73%)	139013 (32.58%)
Hispanic	7102 (15.69%)	92020 (21.57%)
Asian/PI	2049 (4.53%)	6858 (1.61%)
Other	6169 (13.62%)	41278 (9.67%)
Age	Set to Nearest 18-Month Well Child Visit	Set to 18 Months

Comparative Analysis of Model Performances

Figure 1 and Figure 2 both display model performances based on Area Under the Recievier Operator Curve (AUROC or AUC) across all four model types as the sample size of the target data is varied. In both tasks, both EWC and VCL borrowing models demonstrated higher performance in AUC compared to the models on target data only when the sample size was small. Through this result, we observe a significant increase in performance for borrowing models trained on very small samples, compared to a model trained only on Duke data, as long as the performance of external data is higher than the target. That is, the borrowing models learn the appropriate amount of borrowing to increase performance on its target, and serves as safeguards to a model’s performance when target sample size is small.

Figure 1.

AUROC performance metrics of the ASD prediction task, across varied sample sizes of the target data’s training set, with corresponding 95% bootstrap confidence intervals.

Figure 2.

AUROC performance metrics of the ESRD rehospitalization prediction task, across varied sample sizes of the target data’s training set, with corresponding 95% bootstrap confidence intervals.

Comparative Analysis of Model Calibration

Figure 3 and Figure 4 both similarly display the calibration slopes across all models. As the sample size of the target data used to train each model decreases, the calibration slopes of models trained only on target data decrease drastically. However, calibration slopes of the borrowing models stay close to the calibration slope of the model trained on the external data when the target data’s size is made small. Furthermore, the borrowing model using VCL has noticeably superior calibration slope across all models compared to its EWC counterpart in the ASD task, but not the ESRD task. Through this result, we observe a protective effect on the model’s calibration when the sample size is small.

Figure 3.

Calibration slopes of the ASD prediction task, across varied sample sizes of the target data’s training set.

Figure 4.

Calibration slopes of the ESRD rehospitalization prediction task, across varied sample sizes of the target data’s training set.

Analysis of the Degree of External Borrowing

Figure 5 and Figure 6 both characterize the functional similarities between models as sample size is varied. The correlations of predictions made on the fixed test set between the borrowing models and the models on Duke data alone decreased for both tasks as the sample size was made small. But the correlations of predictions made on the fixed test set between the borrowing models and the external data-only models remained high for both tasks. These trends show the borrowing models’ ability to prioritize external data when target data is small. This indicates that the borrowing models functionally behave more similarly to the model on external data when target data is small, providing further proof that it learns the appropriate amount of borrowing.

Figure 5.

Correlation between model predictions on the held-out target data for the ASD prediction task, across varied sample sizes of the target data’s training set.

Figure 6.

Correlation between model predictions on the held-out target data for the ESRD rehospitalization prediction task, across varied sample sizes of the target data’s training set.

Conclusion

The experiments above show that the borrowing models adaptively incorporated information from external data without overpowering the signal from the target data. This suggests that utilizing continual learning methods to borrow external data may impart benefits in performance EHR-based deep learning models. Further, we provide empirical evidence that continual learning approaches dynamically adjust the degree of borrowing, and has good properties when the training sample is small.

Deep learning performs better when trained on large datasets. However, if a small health system seeks to build an EHR-based predictive model and have concerns that the size of their data is too small, they may consider utilizing a continual learning method to borrow external data. We find that these methods safeguard against model performance decline in small samples by borrowing and regularizing parameters from external sources.

Discussion

To our knowledge, our work is the first to suggest utilizing continual learning methods to improve the performance of EHR-based predictive models on small samples by leveraging external data. The closest precedent to the use of continual learning in EHR-based deep learning is the WUPERR (Weight Uncertainty Propagation and Episodic Representation Replay) algorithm, which integrates EWC with episodic replay for sepsis prediction across hospitals³. However, WUPERR focuses on catastrophic forgetting in sequentially added cohorts rather than our adaptation of continual learning to augment models on small samples.

Several strategies have been developed to address limited sample sizes in EHR data, predominantly through data augmentation and synthetic generation^{21, 22}. Synthetic data approaches using generative adversarial networks (GANs) and its derivatives have emerged as prominent solutions in this area²³. For data augmentation, recombination-based methods and permutation-based methods strategically rearrange medical records systematically transform EHR records while still retaining the underlying structure and semantics of the source data²⁴. While these methods effectively expand single datasets through artificial examples or knowledge integration, our continual learning approach fundamentally differs by directly leveraging existing external datasets rather than generating new training instances, enabling cross-institutional knowledge transfer without raw data synthesis.

Utilizing continual learning models for external borrowing simultaneously allows for privacy-preserving EHR-based deep modeling while approaching the problem of having limited sample sizes. Transfer learning methods are also commonly used to address this challenge by using pretrained models on a source before fine-tuning on target data²⁵. However, we opted to focus on continual learning instead as these have an explicit knowledge retention mechanism helpful when modeling with a small target dataset. Optimal transport approaches provide another familiar avenue for privacy-preserving knowledge transfer by aligning feature distributions between source and target populations, but requires access to the distributions of both source and target datasets²⁶. These methods differ from continual learning approaches like EWC and VCL, as they do not require fine tuning weights or distribution alignment to enable external data borrowing.

Despite this, there are some limitations to utilizing continual learning methods in a data borrowing context. EWC and VCL both requires identical network structures between a model on target data and a model on external data. Further, the benefits of EWC and VCL may depend on the choice of external data and its alignment with the target population. The external data that we used for our demonstration were similar in outcome rate as well as the proportion of various levels across demographic factors. However, we have not tested cases where the external data distribution which we borrow diverges significantly from the target data. Moreover, while this study focused on binary classification tasks, extending these methods to multiclass or regression tasks can provide insight on this framework’s applicability to other types of models. Future work could explore strategies for assessing data compatibility to refine the borrowing process. Additionally, investigating how continual learning frameworks interact with different neural architectures and optimization strategies may further enhance their utility in EHR-based predictive modeling.