Title: GRAM: Graph-based Attention Model for Healthcare Representation Learning
Authors: Edward Choi
Mohammad Taha Bahadori
Le Song
Walter F. Stewart
Jimeng Sun
Venue: ICLR 2017
Keywords: Deep learning
Applications

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 models 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.

TL;DR: We propose a novel attention mechanism on graphs to learn representations for medical concepts from both data and medical ontologies to cope with insufficient data volume.

Review discussion:
Hi,
I do not understand well the argument at page 6, where you speculate on the role of initialized embeddings on different tasks.
I cannot understand why the initialized embeddings are more useful on a multiclass task than on a binary one. Can it be that it depends more on the input data that on the nature of the task?
Are the initialized embeddings fine-tuned duing training?
The sequential disease prediction (SDP) task and the heart failure (HF) prediction task have fundamentally different natures. SDP tries to predict the occurrence of  diagnosis codes in the next visit V_{t+1} based on the information of past co-occurring diagnosis codes V_1, ..., V_t, as opposed to predicting a binary label for the entire sequence in HF task. The co-occurrence information infused by the initialized embedding technique significantly improves the performance of SDP, because it is more about finding co-occurrence of the codes including the same codes repeating over time. The dataset we used for the HF prediction task is a subset of the large Sutter PAMF dataset, which has also been used for one of the two SDP tasks (The construction of the HF dataset  is described in Appendix B). The characteristic of the two datasets are quite similar.
Therefore, we believe the benefit of the initialized embedding is more evident in the SDP task because of the differences in the learning tasks rather than the datasets.
As for your last question, yes, the initialized embeddings are fine-tuned during the training, but only for the leaf nodes. We do not fine-tune the non-leaf nodes during the training.
We will clarify both points in the manuscript. Thanks for your questions.
What determined the choice of the specific form of the function computing the compatibility between embeddings e_i and e_j (Eq. 3)?
For instance, why would a simple dot product between e_i and e_j not work?
Thank your for your question. A feedforward neural network with a single hidden layer like our formulation is well known to be a sufficient approximator for an arbitrary function including a dot product.
Moreover, our formulation resembles the compatibility used in the original attention model  Bahdanau et al. 2014, where they calculate the compatibility between the encoder’s hidden state and the decoder’s hidden state with “v_a dot tanh(W_a s_i + U_a h_j)”.
However, after testing both formulations, we empirically found that our formulation performed better in our use cases.
SUMMARY.
This paper presents a method for enriching medical concepts with their parent nodes in an ontology.
The method employs an attention mechanism over the parent nodes of a medical concept to create a richer representation of the concept itself.
The rationale of this is that for  infrequent medical concepts the attention mechanism will rely more on general concepts, higher in the ontology hierarchy, while for frequent ones will focus on the specific concept.
The attention mechanism is trained together with a recurrent neural network and the model accuracy is tested on two tasks.
The first task aims at prediction the diagnosis categories at each time step, while the second task aims at predicting whether or not a heart failure is likely to happen after the T-th step.
Results shows that the proposed model works well in condition of data insufficiency.
OVERALL JUDGMENT
The proposed model is simple but interesting.
The ideas presented are worth to expand but there are also some points where the authors could have done better.
The learning of the representation of concepts in the ontology is a bit naive, for example the authors could have used some kind of knowledge base factorization approach to learn the concepts, or some graph convolutional approach.
I do not see why the the very general factorization methods for knowledge bases do not apply in the case of ontology learning.
I also found strange that the representation of leaves are fine tuned while the inner nodes are not, it is a specific reason to do so?
Regarding the presentation, the paper is clear and the qualitative evaluation is insightful.
DETAILED COMMENTS
Figure 2. Please use the same image format with the same resolution.
We would like to thank you for your questions and insightful comments. We address your comments in detail below.
1. Using the knowledge base factorization approach or graph convolutional approach?
As discussed in the Related Work section, knowledge base factorization methods such as Neural Tensor Network [Socher et al. 2013], TransE [Bordes et al. 2013] or TransH [Wang et al. 2014] learn the representations of entities and their relations in the knowledge graph, in order to perform tasks such as link prediction, triple classification or knowledge completion. Our goal, on the other hand, is fundamentally different in that we do not aim to learn the structure of the ontology (i.e. learning the entities and various relationships between entities), but focus on the parent-child relationship between medical concepts and learn how to properly combine the ancestors in the knowledge DAG to obtain robust representations for medical concepts under data insufficiency.
Graph convolutional approaches such as Planetoid [Yang et al. 2016] or GCN [Kipf and Welling 2016] focus on learning the node representations to mainly perform node classification. Similar to the works mentioned above, they use the graph itself as the main data and perform prediction tasks for that graph. Our work, on the other hand, uses longitudinal EHR as the main data and perform prediction tasks for that EHR, using the knowledge DAG as the supplementary information.
Also, graph convolutional approaches are similar to DeepWalk [Perozzi et al. 2014] or node2vec [Glover and Leskovec 2016] mentioned in the Related Work section in that they focus on graphs with flat hierarchy where all nodes are equal whereas GRAM utilizes hierarchical relationship (i.e. parent-child) of nodes.
We added more works to discuss them in the Related Work section.
2.
Regarding the fine-tuning of the nodes
We apologize for the confusion caused by our previous response. The basic embeddings of both leaf nodes and non-leaf nodes are fine-tuned during the training process, since the error signal flows from the output \widehat{y}_t to the final representations g_i's which are convex combinations of the basic embeddings e_i's. We apologize for the confusion. We clarified this in the manuscript.
3. Figure 2. Please use the same image format with the same resolution.
We replaced Figure 2 with tables for better readability.
Also Figure 6 in the Appendix was replaced with tables as well.
I read the authors' response and maintain my rating.
This paper introduces an approach for integrating a direct acyclic graph structure of the data into word / code embeddings, in order to leverage domain knowledge and thus help train an RNN with scarce data. It is applied to codes of medical visits. Each code is part of an ontology, which can be represented by a DAG, where codes correspond to leaf nodes, and where different codes may share common ancestors (non-leaf nodes) in the DAG.
Instead of embedding merely the leaf nodes, one can also embed the non-leaf nodes, and the embeddings of the code and its ancestors can be combined using a convex sum. That convex sum can be seen as an attention mechanism over the representation. The attention weights depend on the embeddings and the weights of an MLP, meaning that the model can separate learning the code embeddings and the interaction between the codes. Embedding codes are pretrained using GloVe, then fine-tuned.
The model is properly evaluated on two medical datasets, with several variations to isolate the contribution of the DAG (GRAM or GRAM+ vs. RNN or RandomDAG) and of pretraining the embeddings (RNN+ vs RNN, GRAM+ vs GRAM). Both are shown to help achieve the best performance and the evaluation methodology seems thorough.
The paper is also well written, and the case for MLP attention instead of a plain dot product of embeddings was made by the authors.
My only two comments would be:
1) Why is there a softmax in equation 4, given that the loss is multivariate cross-entropy (in the predicted visit, several codes could be equal to 1), not a a single-class cross-entropy?
2) What is the embedding dimension m?
Thank you for your review and clarifying questions. We address your questions in detail below.
1. Why is there a softmax in equation 4, given that the loss is multivariate cross-entropy (in the predicted visit, several codes could be equal to 1), not a single-class cross-entropy?
The loss function in section 2.3 is a binary cross-entropy for each code we are predicting.
It is true that multiple codes could be 1 in a single visit. Considering this, we also experimented with a set of sigmoid functions, one for each code.
But using the Softmax function showed better predictive performance. We believe this is due to the strong regularizing effect caused by the normalizing denominator of Softmax.
Also, although there could be multiple codes equal to 1 in a single visit, the average number of codes per visit is limited to 2 in Sutter dataset and 13 in MIMIC-III dataset, and Softmax seems sufficient to capture that code distribution.
2. What is the embedding dimension m?
The embedding dimension m is the size of the vector we used to represent the nodes (i.e. medical concepts) in the knowledge DAG. The values of m we used in the experiments, along with other hyperparameters, are described in Table 4 in Appendix C.
This paper addresses the problem of data sparsity in the healthcare domain by leveraging hierarchies of medical concepts organized in ontologies. The paper focuses on sequential prediction given a patient’s medical record (a sequence of medical codes, some of which might occur very rarely).
Instead of simply assigning each medical code an independent embedding before feeding it to an RNN, the proposed approach assigns each node in the medical ontology a “basic” embedding, and composes a “final” embedding for each medical code by taking a learned weighted average (via an attention mechanism) of the medical code’s ancestors in the ontology.
Notably, the paper is well written and the approach is quite intuitive.
I have the following comments:
- Why is the patient’s visit taken as just the sum of medical codes found in the visit, and not say the average or a learned weighted average? Wouldn’t this bias for/against the number of codes in the visit?
- I don’t see why basic embeddings are not fine tuned as well. Did you find that to hurt performance? Do you have an explanation for that?
- Looking at Figure 2, the results seem very close and the figures are not very clear (figure (b) top is missing).
Also, I am wondering how significant the differences are so it would be nice to comment on that.
Finally, I think this is an interesting application paper applying well-established deep learning techniques. The paper deals with an important issue that arises when applying deep learning models in domains with scarce data resources.
However, I would like the authors to comment on what there paper offers as new insights to the ICLR community and why they think ICLR is a good avenue for their work.
Thank you for your review and your insightful questions. Our responses to your comments are as below.
1. Why is the patient’s visit taken as just the sum of medical codes found in the visit, and not say the average or a learned weighted average? Wouldn’t this bias for/against the number of codes in the visit?
The sum of medical codes is a natural approach as a patient (or a visit) can be viewed as the aggregate of medical symptoms and diseases, as opposed to averaging, which does not capture the complexity of the patient status. Your suggestion on the learned weighted average is related to one of our future research directions. Applying another attention mechanism to weight-sum the code embeddings could be a better way of representing a visit.
However, under a proper supervised setting, the embedding matrix for the codes would learn values that incorporates the attention weights as well. This could potentially be an interesting comparison and can be a focus for another paper.
2. I don’t see why basic embeddings are not fine tuned as well. Did you find that to hurt performance? Do you have an explanation for that?
We apologize for the confusion. As we explained in the response for the first review, the basic embeddings for both leaf nodes and non-leaf nodes are fine-tuned, since the error signal flows from the output \widehat{y}_t to the final representations of the code g_i’s, which are convex combinations of the basic embeddings e_i’s. We again apologize for the confusion. We clarified this in the paper.
3. Looking at Figure 2, the results seem very close and the figures are not very clear (figure (b) top is missing).
Also, I am wondering how significant the differences are so it would be nice to comment on that.
We replaced Figure 2 with tables for better readability. We also replaced Figure 6 in the Appendix with tables. In sequential diagnoses prediction, especially for MIMIC-III where the training data is even smaller, GRAM+ shows significantly better performance compared to RNN and RNN+ (maximum 6% improved AUC for Sutter PAMF dataset, 10% improved AUC for MIMIC-III) for infrequent codes. SimpleRollUp is quite effective since it is also uses prior knowledge.
For HF prediction, GRAM and GRAM+ consistently outperform other baselines. I believe these improvements are quite significant and hard to achieve given that the strong baseline performance (AUC>.8). We added more descriptions to the paper to emphasize the performance differences.
4.
Finally, I think this is an interesting application paper applying well-established deep learning techniques. The paper deals with an important issue that arises when applying deep learning models in domains with scarce data resources.
However, I would like the authors to comment on what there paper offers as new insights to the ICLR community and why they think ICLR is a good avenue for their work.
Thank you for appreciating the value of this work. To our knowledge, this work is the first attempt to apply the attention mechanism to a graph (DAG to be precise) in order to augment prior knowledge when learning representations. We believe that we have shown this to be not only viable, but also quite effective. We believe our work will invite the exciting follow-up research in both deep learning methods as well as important applications to healthcare to 1) invent different, or even better methods to apply attention on graphs and 2) introduce various graph-based knowledge to learn robust representations in healthcare as well as other fields.
Major changes in the revision
-- Figures 2 are replaced with tables for better readability
-- Related works have been expanded
-- Incorporated the reviewers’ comments in the experiment section.
-- Experiments for GRAM+ and RNN+ in heart failure prediction were re-run
More details
We discovered that, for GRAM+ and RNN+ in HF prediction with varying amounts of training data, we initialized the basic embeddings e_i’s and the embedding matrix W_emb with GloVe vectors that were trained always on the 100% training data instead of the downsampled training data. This could have exaggerated the AUC of both models.
Therefore we conducted a correct experiment where the initialization was done with GloVe vectors trained on the downsampled training data. We found that GRAM+’s performance did not show noticeable difference, but RNN+’s performance dropped approximately 0.5-1% AUC.

Meta-review:
role: PC
data: comment: The reviewers all agreed that the paper was well written, that the proposed approach is very sensible and intuitive and that the experiments are convincing. However, they are concerned that the proposed work is of limited interest to the ICLR community. The technical contribution is not significant enough for any of the reviewers to strongly recommend an acceptance.

Decision: Reject