Source-linked AI summary
MiME: Multilevel Medical Embedding of Electronic Health Records for Predictive Healthcare
Edward Choi, Cao Xiao, Walter F. Stewart, Jimeng Sun
TL;DR
Predictive healthcare models often require more EHR data than many systems have, and external ontologies may not match local terminology. MiME models EHRs’ inherent multilevel code structure while jointly learning from auxiliary prediction tasks, and it outperformed baselines across the evaluated tasks, including a 15% relative PR-AUC gain on the smallest heart-failure dataset.
Problem
Deep learning models for EHR prediction typically require data volumes exceeding most healthcare systems’ capacity, while medical ontologies may be unavailable or unsuitable because of terminology inconsistencies.
Method
MiME learns diagnosis-, visit-, and patient-level embeddings by modeling diagnosis–treatment relationships and jointly predicting structure-based auxiliary targets without external labels.
Results
MiME outperformed baseline methods across heart failure and sequential disease prediction settings, with a 15% relative gain in PR-AUC over the best baseline on the smallest heart-failure dataset.
Takeaways & Limitations
MiME’s inherent-structure modeling was associated with improved predictive performance and generalization on smaller EHR datasets, particularly for heart failure prediction.
Abstract
from arXiv · showhide
Deep learning models exhibit state-of-the-art performance for many predictive healthcare tasks using electronic health records (EHR) data, but these models typically require training data volume that exceeds the capacity of most healthcare systems. External resources such as medical ontologies are used to bridge the data volume constraint, but this approach is often not directly applicable or useful because of inconsistencies with terminology. To solve the data insufficiency challenge, we leverage the inherent multilevel structure of EHR data and, in particular, the encoded relationships among medical codes. We propose Multilevel Medical Embedding (MiME) which learns the multilevel embedding of EHR data while jointly performing auxiliary prediction tasks that rely on this inherent EHR structure without the need for external labels. We conducted two prediction tasks, heart failure prediction and sequential disease prediction, where MiME outperformed baseline methods in diverse evaluation settings. In particular, MiME consistently outperformed all baselines when predicting heart failure on datasets of different volumes, especially demonstrating the greatest performance improvement (15% relative gain in PR-AUC over the best baseline) on the smallest dataset, demonstrating its ability to effectively model the multilevel structure of EHR data.
1 Introduction
Deep learning for EHR prediction often needs more data than health systems possess, while external ontologies may not fit local terminology. MiME instead leverages EHRs’ inherent multilevel structure and auxiliary prediction tasks, improving performance across evaluated settings.
- Most health systems lack the data volume needed to optimize deep learning models, especially for rare conditions and less common services.
- Medical ontologies can address data-volume insufficiency, but idiosyncratic terminology and cross-system mapping make them unavailable or unsuitable in many settings.
- EHR structure forms a hierarchy from patients to visits, diagnosis codes, and associated medication or procedure orders.
- MiME learns multilevel EHR embeddings while jointly performing structure-based auxiliary prediction tasks without external labels.
- 15% relative gain in PR-AUC over the best baseline occurred for heart failure prediction on the smallest dataset.
2 Method
MiME represents EHR visits hierarchically by modeling diagnosis codes together with associated treatments, then aggregates these representations into visit and patient embeddings. Auxiliary predictions of diagnosis and treatment codes regularize this representation learning without additional labels.
- MiME explicitly models diagnosis–treatment relationships within visits instead of flattening all codes into one binary visit vector.
- Each diagnosis object combines one diagnosis code with its associated set of medication or procedure codes.
- Visit embeddings are transformed and mapped across a sequence of visits to form a patient representation for downstream prediction.
- Diagnosis-object embeddings aggregate diagnosis representations with treatment interactions, and multiple objects are aggregated into a visit embedding.
- MiME jointly predicts diagnosis and treatment codes as auxiliary tasks using cross-entropy losses weighted by λaux.
- These auxiliary tasks require no additional labeling effort and encourage embeddings that capture visit events while remaining useful beyond the specific target task.
3 Experiments
Experiments evaluated MiME for heart failure prediction across datasets varying in size and visit complexity, using GRU-based patient representations and comparisons with baseline embedding methods. MiME consistently performed strongly, with the largest relative gain on the smallest dataset and improved performance over GRAM without external domain knowledge.
- Datasets and evaluation: The experiments used Sutter Health EHR data from 30,764 senior patients to predict first heart failure diagnosis from an 18-month observation period.The dataset included diagnosis, medication, and procedure codes; 3,414 patients were cases and 27,350 were controls.
- Baselines and settings: MiME was compared with raw, linear, nonlinear, and multilayer embedding strategies, as well as Med2Vec and GRAM, with GRU mapping visit sequences to patient representations.The experiments also varied maximum sequence length to create datasets representing different amounts of collected patient history.
- Datasets and evaluation: PR-AUC was the primary reported metric because the heart failure prediction data were imbalanced.Models were trained end-to-end except Med2Vec, and logistic regression converted patient representations into heart failure probabilities.
- Heart failure prediction: 15% relative improvement: on the smallest E1 dataset, MiME aux achieved 0.2831 PR-AUC versus 0.2462 for the best baseline tanhmlp.MiME outperformed all baseline models across datasets, with larger performance gaps on E1 and E2 than on E3 and E4.
- Comparison with external knowledge: MiME consistently outperformed GRAM in test loss and test PR-AUC while using the inherent EHR structure rather than external treatment-code domain knowledge.GRAM used only the ICD9 diagnosis-code hierarchy and did not include additional treatment-code knowledge.
- Visit complexity: In visit-complexity experiments, MiME and MiME aux were close to the best performers across D1, D2, and D3, while MiME aux outperformed MiME on every reported measure.On high-complexity D3, MiME aux reached 0.4787 PR-AUC, an 8.4% relative improvement over the best baseline tanh.
4 Related Work
Prior medical concept embedding work compressed sparse concepts or represented increasingly abstract concepts, while newer EHR methods considered multiple code types. MiME is positioned against approaches that either omit explicit EHR hierarchy or flatten diagnosis and medication codes, and it also draws on auxiliary-task training.
- Medical concept embeddings: Earlier approaches learned embeddings from temporal sequences of medical concepts or from latent layers of deep models representing more abstract concepts.Examples included restricted Boltzmann machines, stacked auto-encoders, and multilayer neural networks.
- Multiple EHR code types: Some multi-code EHR methods modeled disparate code types with RNNs but did not explicitly address the hierarchy of EHR data.Other work modeled interactions between diagnosis and medication-code sets, yet simple concatenation often outperformed those interaction methods.
- Auxiliary-task learning: Prior medical prediction studies used related auxiliary targets, including laboratory values, mortality, and phenotyping, to improve model robustness or prediction accuracy.MiME’s auxiliary-task strategy belongs to this broader line of multitask medical prediction research.
5 Conclusion
MiME jointly models hierarchical inter-code relations and auxiliary prediction tasks, showing strong benchmark performance and generalization to smaller datasets. The authors identify future extensions to finer-grained medical events.
- MiME simultaneously models hierarchical inter-code relations in medical concept embeddings and performs auxiliary prediction tasks.
- MiME demonstrated performance across benchmark tasks and generalized to smaller datasets, especially for heart failure prediction measured by PR-AUC.
- Future work will extend MiME to procedure outcomes, demographic information, and medication instructions.
- The heart failure case-control selection used patients aged 40 to 85 and matched each case with up to ten eligible controls.
- Table 5 reports heart failure prediction performance on small datasets, with standard deviations from 5-fold random data splits and the two best values in each column highlighted.
C Training Details
The experiments used fixed TensorFlow and optimization settings, repeated random data splits, and model-specific representation sizes selected to support fair comparisons across dataset scales.
- All models used TensorFlow 1.4, Adam optimization, and a learning rate of 1e-3.
- Results were averaged over 5-fold random splits using 70% training, 10% validation, and 20% test data.
- Models used minibatches of 20 patients for 20,000 iterations, with validation-loss evaluation every 100 iterations for early stopping.
- For small datasets, baseline visit vectors were generally 128-dimensional, while MiME embeddings were adjusted to match baseline parameter counts.
- For progressively larger datasets, baseline visit vectors and GRU cell sizes were set to 256, with corresponding MiME adjustments.
D Heart Failure Prediction Performance on Datasets D1, D2 and D3, Full Version
On datasets D1, D2, and D3, bounded sigmoid and tanh activations generally outperformed ReLU, although sigmoid’s regularization benefit diminished with more data.
- Sigmoid and tanh mostly outperformed ReLU on both measures across D1, D2, and D3.
- The bounded outputs of sigmoid and tanh may regularize summed code embeddings, with the effect especially favorable in D3.
- The regularization benefit of sigmoid and tanh diminished as the dataset grew.
- Sigmoidmlp showed evidence of vanishing-gradient behavior compared with tanh and ReLU, which have larger gradient values.
E ROC-AUC of Heart Failure Prediction on Datasets D1, D2 and D3
MiME generally outperformed baseline models in ROC-AUC on the small heart failure datasets, but its advantage was smaller than for PR-AUC. The imbalance makes specificity easier than identifying positive cases.
- Table 6 reports ROC-AUC for all models on small datasets, including standard deviations from 5-fold random splits and two highlighted best values per column.
- Except on D1, where visit complexity was low, MiME consistently outperformed all baseline models in ROC-AUC.
- The ROC-AUC gap between MiME and baselines was smaller than the corresponding PR-AUC gap.
- PR-AUC more strongly tests retrieval of positive cases while ignoring negatives, which is harder in problems with few positive samples.
- Heart failure prediction contains many more controls than cases, making high specificity relatively easy to achieve.
F Test PR-AUC on Datasets E1, E2, E3 and E4, Full Version
Test PR-AUC varied across dataset sizes and baseline models, with different baselines strongest in small versus large datasets.
- Baseline PR-AUC performance fluctuated as dataset size increased.tanhmlp performed competitively on small datasets but weakened on large datasets, whereas relumlp became the best baseline on large datasets.
- tanhmlp was competitive on small datasets but weaker on large datasets.
- relumlp did not stand out on small datasets but became the best baseline on large datasets.
G Test Loss and Test ROC-AUC on Datasets E1, E2, E3 and E4
Across datasets E1–E4, MiME and MiME aux generally outperformed the baselines on test loss and test ROC-AUC, with MiME aux usually performing best.
- MiME and MiME aux consistently outperformed all baselines on test loss and test ROC-AUC, except Med2Vec.
- MiME aux outperformed MiME on both metrics except test loss in E4.
- MiME aux showed its strongest advantage over MiME on the smallest dataset E1.
- Sequential disease prediction uses past visits to predict diagnosis codes in the next visit.The inputs include diagnosis and treatment codes, while outputs contain only diagnosis codes.
I Experiment Results for Sequential Disease Prediction
MiME performed best on sequential disease prediction across reported metrics and generally improved prediction for diagnosis codes outside the rarest group.
- MiME achieved the best performance in all reported sequential disease prediction metrics.The evaluation used five-fold random data splits, with standard deviations reported in parentheses.
- MiME’s results support capturing temporal progression of patient status in sequential disease prediction.The task predicts diagnoses in the next visit from all past visits.
- MiME outperformed all baseline models except for the rarest diagnosis-code group.The largest reported improvement was an 11.6% relative gain over tanhmlp.
- Med2Vec performed best for the rarest diagnosis-code group.
- MiME showed good performance across diverse prediction settings and significantly outperformed baselines on the more complex heart-failure prediction task.The heart-failure task involved relationships between labels and codes that were not straightforward.