Source-linked AI summary
GRAM: Graph-based Attention Model for Healthcare Representation Learning
Edward Choi, Mohammad Taha Bahadori, Le Song, Walter F. Stewart, Jimeng Sun
TL;DR
Healthcare deep learning faces data insufficiency and the need for representations aligned with medical knowledge. GRAM uses ontology structure and attention to combine medical concepts with their ancestors, improving prediction performance across sequential diagnosis and heart-failure tasks while producing ontology-aligned representations.
Problem
Healthcare deep learning predictive models face insufficient data and a need for representations aligned with medical knowledge.
Method
GRAM uses a medical knowledge DAG and attention to combine each concept with weighted ancestor information in an end-to-end predictive model.
Results
GRAM improved prediction for data-insufficient diagnoses, achieved 3∼4% AUC gains over most baselines for heart-failure prediction, and learned representations aligned with the ontology.
Takeaways & Limitations
GRAM adaptively gives more weight to ancestors for less-observed concepts, providing coarse-grained information while retaining medically structured representations.
Takeaways & Limitations
On smaller datasets, model scatterplots were less structured, limiting the qualitative evaluation of learned representations.
Abstract
from arXiv · showhide
Deep learning methods exhibit promising performance for predictive modeling in healthcare, but two important challenges remain: -Data insufficiency:Often in healthcare predictive modeling, the sample size is insufficient for deep learning methods to achieve satisfactory results. -Interpretation:The representations learned by deep learning methods should align with medical knowledge. To address these challenges, we propose a GRaph-based Attention Model, GRAM that supplements electronic health records (EHR) with hierarchical information inherent to medical ontologies. Based on the data volume and the ontology structure, GRAM represents a medical concept as a combination of its ancestors in the ontology via an attention mechanism. We compared predictive performance (i.e. accuracy, data needs, interpretability) of GRAM to various methods including the recurrent neural network (RNN) in two sequential diagnoses prediction tasks and one heart failure prediction task. Compared to the basic RNN, GRAM achieved 10% higher accuracy for predicting diseases rarely observed in the training data and 3% improved area under the ROC curve for predicting heart failure using an order of magnitude less training data. Additionally, unlike other methods, the medical concept representations learned by GRAM are well aligned with the medical ontology. Finally, GRAM exhibits intuitive attention behaviors by adaptively generalizing to higher level concepts when facing data insufficiency at the lower level concepts.
1 Introduction
Healthcare predictive modeling increasingly uses deep learning with expanding EHR data, while GRAM represents medical concepts through ontology-based ancestor information before prediction.
- Deep learning methods are increasingly applied to predictive modeling with rapidly growing and diverse healthcare data.
- GRAM combines leaf-concept embeddings with ancestor embeddings through attention to form concept representations.These representations form an embedding matrix used to embed patient visits before neural-network prediction.
- GRAM transforms patient visit vectors into visit representations that a neural network uses to predict target outcomes.
2 Methodology
GRAM uses medical ontologies to enrich EHR representations, combining each concept with its ancestors through learned attention. These representations feed an end-to-end predictive model, with ancestor embeddings initialized from augmented co-occurrence statistics.
- Knowledge DAG: GRAM represents medical concepts using parent-child relationships in a medical ontology, whose ancestors provide progressively more general concepts.The ontology is modeled as a directed acyclic graph, with EHR codes as leaf nodes and internal nodes as ancestors.
- Attention mechanism: When data are constrained, GRAM adaptively increases ancestor influence because higher-level concepts can be learned more accurately and provide coarse-grained information.An attention mechanism automates this shift in specificity during end-to-end training.
- Attention mechanism: Each leaf code’s final representation is a convex combination of its own and its ancestors’ basic embeddings, weighted by learned attention coefficients.The attention weights are computed for the code and all of its ancestors, and the resulting representations form the code embedding matrix.
- Attention mechanism: GRAM computes ancestor attention with an MLP compatibility function that concatenates child and ancestor embeddings in that order.The paper reports that this MLP formulation performed better in its use cases than inner-product and Bahdanau-style alternatives.
- End-to-end predictive modeling: Final code representations are combined within each visit and passed through an RNN to predict subsequent diagnoses or other target labels.The visit representation is obtained by multiplying the embedding matrix by a multi-hot visit vector, after which the RNN produces predictions and a cross-entropy loss.
- Embedding initialization: GRAM initializes leaf and ancestor embeddings with GloVe trained on co-occurrence statistics from visits augmented with each code’s ancestors.Higher ancestors appear more frequently in augmented visits, causing general concepts to have greater involvement in co-occurrence statistics; embeddings are then fine-tuned during model training.
3 Experiments
GRAM was evaluated on sequential diagnosis and heart-failure prediction tasks against recurrent and ontology-based baselines, with analyses of data insufficiency, scalability, and interpretability. It improved performance for infrequent diagnoses and heart-failure prediction while producing representations aligned with medical ontology structure.
- Experiment Setup: The experiments covered two sequential diagnoses prediction datasets and one heart-failure onset prediction dataset.The datasets included Sutter PAMF, MIMIC-III, and a Sutter heart-failure cohort.
- Experiment Setup: GRAM, GRAM+, RNN, RNN+, RandomDAG, SimpleRollUp, and RollUpRare provided comparisons across ontology, initialization, and infrequent-code handling strategies.GRAM used ontology-based attention, while GRAM+ additionally initialized basic embeddings using the scheme described in Section 2.4.
- Prediction Performance: 10% higher accuracy was achieved by GRAM+ than the basic RNN in MIMIC-III’s 20th–40th percentile range of diagnosis frequencies.The gain occurred where labels were less observed in training, and was greater on MIMIC-III than on Sutter data.
- Prediction Performance: 3–4% AUC improvement over most baselines was achieved by GRAM and GRAM+ for heart-failure prediction, while exceeding RNN+ by a maximum of 1.8% AUC.The AUC values were already in the mid-80s; GRAM+ provided limited initialization benefit over GRAM for this task.
- Qualitative Evaluation: GRAM and GRAM+ learned disease representations that were more structured and more consistent with the medical knowledge DAG than the other compared representations.The t-SNE evaluation compared GRAM variants with RNN variants, RandomDAG, GloVe, and Skip-gram; the MIMIC-III and Sutter heart-failure plots were less structured because of smaller data size.
- Attention Behavior: GRAM assigned more attention to ancestors for rarely observed diseases, while frequently observed diseases received stronger attention at the leaf level.For example, rare other pneumothorax received most information from its highest ancestor, whereas frequent essential hypertension received strong leaf attention.
4 Related Work
Prior work has explored attention, graph representation learning, knowledge-graph embeddings, and graph regularization, but GRAM focuses on ontology-based attention for medically interpretable prediction under data insufficiency.
- Attention mechanisms have been applied in speech recognition, computer vision, and healthcare, but not previously designed around knowledge ontologies in this work’s framing.
- Graph representation methods learn vertex representations from neighbor information using random walks, breadth-first search, or graph convolutions.
- Knowledge-graph embedding studies project entities and relations into latent spaces for link prediction, triple classification, or entity classification.
- GRAM differs by using the knowledge DAG as a prior for intuitive attention, addressing data insufficiency while learning medically interpretable representations.
- Graph Laplacian regularization incorporates side information, but requires an appropriate graph-distance definition that may be unavailable.
5 Conclusion
GRAM addresses data insufficiency in healthcare prediction by combining EHR data with medical-ontology structure through end-to-end graph-based attention. Experiments report improved prediction performance, especially for low-frequency diseases and small datasets, while attention analysis provides intuitive insight.
- GRAM combines a knowledge DAG with EHR data to learn accurate and interpretable medical-concept representations.
- GRAM represents each medical concept as a weighted average of its ancestors and trains attention jointly with prediction end to end.
- Three real-EHR predictive modeling experiments showed significant performance improvements, especially for low-frequency diseases and small datasets.
- Attention-behavior analysis provided intuitive insight into how GRAM uses ontology structure during prediction.