[Paper Review] GRAM: Graph-based Attention Model for Healthcare Representation Learning

Edward Choi, Mohammad Taha Bahadori+3 more|arXiv (Cornell University)|Nov 21, 2016

Machine Learning in Healthcare28 references99 citations

TL;DR

GRAM integrates medical ontologies into a graph-based attention mechanism to form robust, interpretable leaf-code representations for EHR data, improving predictive performance under data scarcity.

ABSTRACT

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.

Motivation & Objective

• Motivate deep learning for healthcare prediction under data insufficiency and the need for interpretable representations aligned with medical knowledge.
• Leverage hierarchical medical ontologies to enrich leaf medical concept representations.
• Develop an end-to-end trainable attention mechanism that combines leaf concepts with ancestors based on data-driven compatibility scores.

Proposed method

• Represent leaf medical codes as a convex combination of their embeddings and those of their ancestors using an attention mechanism.
• Define a compatibility function f(ei, ej) via a 1-hidden-layer MLP to compute attention weights over ancestors.
• Train the attention module end-to-end with a predictive neural network (RNN/GRU/GRU variants) for disease onset tasks.
• Initialize basic embeddings of leaf and non-leaf nodes using GloVe-derived co-occurrence embeddings over augmented visits that include ancestors.
• Embed visits by multiplying the leaf-and-ancestor embedding matrix with multi-hot visit vectors and feed into an RNN for prediction.
• Compare GRAM variants to RNN baselines and ablations across sequential diagnoses prediction and heart failure onset tasks.

Experimental results

Research questions

• RQ1 Can GRAM improve predictive accuracy for diseases that are rare in the training data compared to standard RNNs?
• RQ2 Do ontology-informed leaf representations align with medical knowledge, and are they interpretable through attention patterns?
• RQ3 How does GRAM perform under varying data availability and across diverse healthcare datasets (Sutter PAMF, MIMIC-III, Sutter HF cohort)?
• RQ4 What is the impact of initializing leaf/non-leaf embeddings with ontology-aware co-occurrence (GloVe) on downstream tasks?

Key findings

• GRAM achieves up to 10% higher accuracy than basic RNN for rare diseases.
• GRAM improves AUC by up to 3% for heart failure prediction with markedly less training data.
• Ontology-informed leaf representations align with medical knowledge and exhibit interpretable attention, adapting to data availability by weighting ancestors more when data are sparse.
• Initialization of basic embeddings with ontology-aware co-occurrence (GRAM+) yields similar or improved performance, especially for sequential diagnosis prediction.
• GRAM adds modest training overhead (~50% more time per epoch) but yields better performance under data scarcity.