Feature Restricted Group Dropout for Robust Electronic Health Record Predictions
Abstract
Recurrent neural networks are commonly applied to electronic health records to capture complex relationships and model clinically relevant outcomes.
However, it is commonplace for the covariates in electronic health records to change distributions.
This work extends restricted feature interactions in recurrent neural networks to address foreseeable and unexpected covariate shifts.
We extend on the previous work by 1) Introducing a deterministic feature rotation so that hyperparameter tuning can search through all combinations of features, 2) Introduce a sub-network specific dropout to ablate the influence of entire features at output of the hidden network, and 3) Extend the feature restrictions to the GRU-D network, which has been shown to be a stronger baseline for covariate shift recovery.
We show that feature restricted GRU-D’s may be more robust to certain perturbations.
Manual intervention was not needed to confer robustness.
Despite this, the LSTM was still the best model in nearly 50% of the cases.
Introduction
In healthcare, it is inevitable that the generation of the data reflected in electronic health records (EHRs) will evolve; new equipment is added with different margins of error, suppliers of laboratory tests will switch and clinicians administering the tests change.
Under these anticipated model changes, how can we be confident that the deep learning model we have deployed can continued to be used until it is remediated or replaced?
We define this problem as one of a sudden covariate change, p ( x i ) →
q ( x i ) .
sudden covariate change is distinct from continual learning setting as the covariate changes rapidly.
Sudden covariate shift is more mild than complete domain shift p ( X ) → q ( X ) .
This makes it challenging, because it raises doubts as to whether or not a model should be taken offline.
If a model has been demonstrated to improve care, and the covariate shift is marginal, it may be unnecessary or even harmful to withhold the predictions.
We also permit p ( y | x i ) ̸
= q ( y | x i ) , such that even if the old covariate distribution is able to be mapped to the new distribution, the patterns in which they are used may shift to reflect healthcare usage.
Such patterns might be present if a new covariate differs in cost, complexity, or speed.
Current methods to mitigate covariate shift often focus on gradual shift.
Work has been done to causally predict which feature will be robust during test time [1].
However this setting assumes that features are already experiencing distributional anomalies.
Once a small subset of the new distribution are available, it is possible continually fine-tune the model, however this could lead to catastrophic forgetting [2].
It is also possible to initiating the model from scratch and upweight the samples in the new distribution [3], however this comes at a training cost and still requires that the model be taken offline while the new distribution is being collected.
Others have sought to apply batch normalisation during inference [4], or Bayesian neural networks
Previous work has shown that restricting the feature interactions in the latent space improves model performance when generalising to an unseen test set [5].
We extend this method by adding a seeded rotation to permit all features to interact with each other.
Crucially, we train the model with a sub-network specific dropout such that each sub-network containing feature x i can be simultaneously omitted.
Finally, we introduce the masking method from Zhang et al.
[5] into the GRU-D model [6], which has shown to be more robust to dramatic distributional changes than LSTM networks [7].
Methods
Data is Sourced from GEMINI , a data warehouse containing 214837 General Internal Medicine inpatient encounters from 6 different sites [8].
The data is extracted and prepared by replicating protocols from previous EHR extractions [9].
Site characteristics are shown in Table 1.
The target selected to demonstrate this task is a delirium label with 30% positively rate.
The data are split by time such that 80% of the training data falls before the training end date, 10% of the data falls between the training end date, and the validation end date, and the remaining 10% occurs after the validation end date and is relegated to the test set.
A GRU-D architecture [6] embeds the site-specific model, while a LSTM network decodes to the target.
The parameters for each site, including the traditional dropout probability, the number of feature interaction groups, the network widths, and the number of decoder layers are selected by taking the network with the best AUROC on a prospective validation set.
Both the unperturbed, feature-restricted GRU-D and the standard GRU-D models are allowed to choose their optimal architectures independently.
These models are compared to an LSTM architecture which has no method to deal with faulty covariates.
To evaluate the robustness of the model, data are randomly perturbed by randomly resampling the feature, flipping the sign of the feature, or by linearly scaling the feature.
We show the performance of the LSTM and GRU-D under these conditions.
As a remedy, we intervene by consciously decaying the input unit of this feature in the GRU-D to the training mean.
The feature-restricted GRU-D is also evaluated on the perturbed input.
Finally, the sub-networks involving that feature are "turned off" by intervening with a sub-network specific dropout.
[Table]
Site
[Table]
[Table]
[Table]
UHNTG  0.
[Table] UHNTW  0. [Table]  Results
The performance of the LSTM, GRU-D, and feature restricted GRU-D are demonstrated in Table 2.
The validation performance under the GRU-D restricted feature interaction tends to outperform
GRU-D and LSTM.
However the generalisation to the prospective test set tends to be weaker with several notable drops in performance from validation time to test time.
This could be attributed to the sites having few limited numbers of samples to do model selection, or to form robust measurements of the AUROC at both validation time and test time.
Alternatively, some of the performance drop can be attributed to real world data-set shift, such as with the introduction of COVID-19 into the hospital system during test time (but not during training or validation time) [Table]
To evaluate the robustness of the model to sudden incremental covariate shifts, one feature is perturbed at a time.
The test set performance of the GRU-D, intervened-upon GRU-D, uncontrolled feature restricted GRU-D, and intervened-upon feature-restricted GRU-D are demonstrated in [-Tables 
 [Table] &-]Tables& 4& 5 for flipped, resampled, and linearly scaled perturbations.
Here we see that both the GRU-D and GRU-D Restricted are capable of handling feature corruption as well as manual intervention.
The corruption itself, however is occasionally devastating, such as the case with UHNTG GRU-D
Restricted model dropping to near random performance with any type of corruption.
[Table]
In general, When the LSTM was the top performer, it lost 0 .
18 · 0 . 24 , 95% CI ) %, 0 .
21 · 0 . 24 , 95% CI ) % and 0 .
14( − 0 .
02 · 0 . 24 , 95% CI ) % of its performance for flipping, resampling, and scaling perturbations, respectively.
In contrast, GRU-D lost it lost 0 .
14( − 0 .
02 · 0 . 43 , 95% CI ) %, 0 .
12( − 0 .
42 , 95% CI ) % and 0 .
32( − 0 .
02 · 0 . 56 , 95% CI ) % of its performance for flipping, resampling, and scaling perturbations, respectively, when it was the top performer.
Compare this to the GRU-D with restriction which lost 0 .
13( − 0 .
10 · 0 . 36 , 95% CI ) %, 0 .
15( − 0 .
05 · 0 . 35 , 95% CI ) % and 0 .
15( − 0 .
18 · 0 . 34 , 95% CI ) % of its performance for flipping, resampling, and scaling perturbations, respectively.
GRU-D was the most susceptible model to scaling noise, while the LSTM was the most susceptible model to flipping and resampling.
Neither intervention taken on the GRU-D or GRU-D Restricted models improved the performance.
After corruption, the LSTM would have still been the best choice 48% of the time, the GRU-D 34% of the time, and the GRU-D with restriction 19% of the time.
[Table]
[Table]
Conclusion
The modification of the GRU-D with the feature restriction mask allows for sub-unit specific off switches.
Across 3 different types of perturbations, the GRU-D with feature restriction experiences less (though not significantly less) harm in performance.
The sub-unit specific off switches confer innate robustness to distribution changes which do not require conscious intervention to utilise.
The robustness of the intervention exceeds that of a GRU-D and LSTM baseline models.
Despite this, the
LSTM architecture was still the preferred architecture in almost half of the cases.
Future work may investigate if sub-unit specific isolation leads to faster fine-tuning or fewer required samples to refit the new sub-unit.