[Paper Review] Learning the Graphical Structure of Electronic Health Records with Graph Convolutional Transformer

Edward Choi, Zhen Xu+5 more|arXiv (Cornell University)|Jun 11, 2019

Machine Learning in Healthcare28 references4 citations

TL;DR

This paper proposes Graph Convolutional Transformer (GCT), a novel model that jointly learns the hidden graphical structure of electronic health records (EHR) and performs supervised prediction tasks, even when structure information is missing. By guiding self-attention with conditional probability priors and attention masks, GCT outperforms baseline models in graph reconstruction and readmission prediction on both synthetic and real-world EHR data.

ABSTRACT

Effective modeling of electronic health records (EHR) is rapidly becoming an important topic in both academia and industry. A recent study showed that using the graphical structure underlying EHR data (e.g. relationship between diagnoses and treatments) improves the performance of prediction tasks such as heart failure prediction. However, EHR data do not always contain complete structure information. Moreover, when it comes to claims data, structure information is completely unavailable to begin with. Under such circumstances, can we still do better than just treating EHR data as a flat-structured bag-of-features? In this paper, we study the possibility of jointly learning the hidden structure of EHR while performing supervised prediction tasks on EHR data. Specifically, we discuss that Transformer is a suitable basis model to learn the hidden EHR structure, and propose Graph Convolutional Transformer, which uses data statistics to guide the structure learning process. The proposed model consistently outperformed previous approaches empirically, on both synthetic data and publicly available EHR data, for various prediction tasks such as graph reconstruction and readmission prediction, indicating that it can serve as an effective general-purpose representation learning algorithm for EHR data.

Motivation & Objective

• To address the challenge of learning effective representations from EHR data when complete structural information (e.g., diagnosis-treatment relationships) is unavailable.
• To investigate whether self-attention in Transformers can be guided to discover meaningful EHR graphical structures without explicit supervision.
• To develop a general-purpose representation learning framework for EHR data that jointly learns structure and performs downstream prediction tasks.
• To empirically validate that incorporating prior knowledge via attention masks and regularization improves structure learning and prediction performance.

Proposed method

• Proposes Graph Convolutional Transformer (GCT), a modified Transformer architecture that integrates prior knowledge into self-attention via attention masks and conditional probability-based regularization.
• Uses the conditional probability matrix P as a prior to guide attention learning, ensuring attention heads focus on plausible relationships between EHR features.
• Applies normalization and learnable parameters to the attention mechanism to balance prior knowledge and data-driven attention patterns.
• Employs a multi-head self-attention mechanism where each head is constrained by a mask derived from the prior probability matrix P.
• Trains the model end-to-end on EHR data using supervised tasks such as graph reconstruction and diagnosis-treatment classification.
• Uses KL divergence and attention visualization to evaluate how well the model’s attention approximates the true adjacency structure.

Experimental results

Research questions

• RQ1 Can self-attention in a Transformer model learn meaningful EHR graphical structures when no explicit structure is provided?
• RQ2 Does incorporating prior knowledge via attention masks and conditional probability regularization improve structure learning and prediction performance?
• RQ3 How does GCT compare to standard Transformers and GCN-based models in learning hidden EHR structures on synthetic and real EHR data?
• RQ4 To what extent does guided attention improve performance on downstream tasks like readmission prediction and graph reconstruction?

Key findings

• GCT outperformed baseline models, including standard Transformers and GCN with fixed priors, on graph reconstruction and readmission prediction tasks across both synthetic and real EHR data.
• GCT achieved lower KL divergence (0.48 on average) between predicted and true adjacency matrices than standard Transformers (0.62), indicating better structure recovery.
• GCT learned attention distributions that closely matched the true structure, especially in early self-attention blocks, and adapted toward the true adjacency over deeper layers.
• The model demonstrated improved performance over GCN with fixed prior (GCN P) in graph reconstruction and diagnosis-treatment classification, showing the benefit of end-to-end structure learning.
• Attention visualization confirmed that GCT initially relied on the prior matrix P but gradually refined attention to align with the true structure, unlike unguided Transformers.
• Even when attention maps were not perfectly aligned with the true structure, GCT still achieved strong performance in graph reconstruction, highlighting the robustness of attention mechanisms in structure discovery.