Source-linked AI summary
Learning the Graphical Structure of Electronic Health Records with Graph Convolutional Transformer
Edward Choi, Zhen Xu, Yujia Li, Michael W. Dusenberry, Gerardo Flores, Yuan Xue, Andrew M. Dai
TL;DR
EHR and claims data may lack the clinical relationships that structure-aware models need, raising whether hidden structure can outperform flat feature representations. The paper proposes GCT, which jointly learns encounter structure and supervised tasks using a guided Transformer, and reports stronger performance across synthetic and public EHR evaluations. Its scope is limited by datasets such as eICU lacking explicit treatment-cause relationships.
Problem
EHR and claims data may lack complete or any explicit structure, so the paper asks whether hidden encounter relationships can be learned instead of treating records as flat bags of features.
Method
GCT jointly learns hidden encounter structure and supervised prediction using Transformer self-attention guided by attention masks and prior conditional-probability regularization.
Results
GCT outperformed baseline models across encounter-based prediction tasks on synthetic data and a publicly available EHR dataset.
Takeaways & Limitations
GCT shows potential as a general-purpose EHR representation-learning algorithm when encounter structure is unknown.
Takeaways & Limitations
The eICU dataset lacks explicit structure, so co-occurring diagnoses and treatments cannot be linked by prescribing reason.
Abstract
from arXiv · showhide
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.
1 Introduction
EHRs contain graphical relationships that reflect clinical decision-making, but common bag-of-features models discard them. GCT addresses missing structure by jointly learning hidden encounter relationships during supervised prediction.
- EHR data include relationships among feature types that reflect the physician’s decision process.Treating an encounter as an unordered feature set discards these relationships, such as why a treatment was ordered.
- Bag-of-features processing loses clinically meaningful distinctions between relationships among diagnoses and treatments.For example, it cannot preserve that Benzonatate was ordered because of Cough rather than Abdominal pain.
- Prior structure-aware modeling improved prediction, but it required EHR data containing encounter structure.This motivates asking whether hidden structure can still improve representations when explicit relationships are unavailable.
- GCT jointly learns hidden encounter structure and supervised prediction tasks when structure information is unavailable.Its proposed Transformer modification uses attention masks and prior conditional-probability regularization to guide structure learning.
- GCT outperformed baseline models across graph reconstruction and readmission prediction on synthetic and publicly available EHR data.The authors present this as evidence of potential for general-purpose EHR representation learning.
2 Related Work
Related work has explored embeddings for patients, visits, and codes, as well as Transformer and graph-attention models. GCT differs by targeting hidden EHR structure jointly with supervised prediction when explicit structure is missing.
- Medical EHR representations: Prior medical concept-embedding work focused on patients, visits, or codes, while the graphical nature of EHR remained incompletely explored.MiME instead derived visit representations according to known encounter structure.
- Transformer models: Existing medical-record Transformer studies mainly replaced recurrent networks or applied BERT objectives without using EHR structure.These approaches addressed sequence processing or representation pretraining rather than explicit hidden-structure learning.
- Transformer and graph networks: Graph-attention research connected self-attention with graph networks by learning edge weights over an adjacency matrix or non-local dependencies.These studies provide context for GCT’s use of self-attention, but do not target the same EHR structure-learning objective.
- Transformer and graph networks: GCT jointly learns underlying EHR structure without explicit structure information while improving supervised prediction tasks.The paper positions this combination as a general-purpose EHR embedding objective.
3 Method
The method represents an encounter as a graph when structure is known, or uses Transformer-based attention to learn hidden connections when it is unavailable. Graph Convolutional Transformer guides this learning with EHR constraints and conditional probabilities while allowing task-useful novel connections.
- Electronic Health Records as a Graph: An encounter is modeled as a graph whose diagnosis, treatment, and lab-code nodes are connected by adjacency matrix A when structure is available.When A is unavailable, feed-forward aggregation effectively treats the encounter as a bag of node representations.
- Transformer and Graph Networks: Transformer learns connection strengths through self-attention by starting from fully connected encounter nodes and refining attention toward meaningful relationships.Unlike sequence processing, positional encoding is unnecessary because encounter features are unordered.
- Graph Convolutional Transformer: Graph Convolutional Transformer restricts attention using a mask that assigns negative infinity to disallowed connections and zero to allowed connections.For example, treatment codes can connect to diagnosis codes but not to other treatment codes.
- Graph Convolutional Transformer: Conditional probabilities such as p(EKG|chest pain) guide likely feature connections and form a row-normalized matrix P aligned with the attention and mask dimensions.The matrix includes probabilities derived from encounter records for diagnosis, treatment, and lab-code relationships.
- Graph Convolutional Transformer: GCT guides attention toward likely meaningful structures while preserving room to learn novel connections useful for supervised prediction tasks.The approach is motivated by recovering the true graph structure without preventing task-relevant deviations.
- Graph Convolutional Transformer: GCT replaces the first Transformer block’s attention with conditional probabilities P, then uses masked self-attention to gradually improve upon that guidance.This design addresses preliminary observations that Transformer attention was often uniformly distributed in its first block.
4 Experiments
Experiments evaluate GCT on synthetic and eICU data using structure reconstruction and prediction tasks. GCT generally performs best, while learned structure and task dependence vary across evaluations.
- Datasets: Synthetic data reproduce the EHR visit-diagnosis-treatment-lab hierarchy, while eICU records lack links connecting treatments to their associated diagnoses.The synthetic data use Pareto-distributed conditional probabilities; eICU contains co-occurring codes but not their causal or clinical relationships.
- Models: GCN uses the true adjacency matrix, whereas GCNP substitutes conditional probabilities and GCNrandom uses a randomly generated normalized adjacency matrix.Shallow and Deep provide non-graph feedforward baselines, with Deep adding layers on top of Shallow.
- Prediction Performance: GCT shows the best performance besides GCN on graph reconstruction and again leads diagnosis-treatment classification; GCNP is outperformed only by GCT.GCNrandom beats Shallow for graph reconstruction but performs worst on diagnosis-treatment classification, indicating task-dependent effects of attending to other features.
- Prediction Performance: On masked diagnosis prediction, GCN is best on synthetic data, while Transformer performs comparably to GCNrandom; eICU scores are higher because its skewed diagnosis distribution makes masking easier.More than 80% of eICU encounters contain diagnosis codes related to prior operating-room presence.
- Prediction Performance: For eICU readmission and mortality prediction, GCT outperforms the other models, although mortality appears to rely little on encounter structure.The authors interpret readmission as benefiting from latent encounter structure, while GCT’s marginal advantage for mortality suggests a broader modeling role.
- Evaluating the Learned Encounter Structure: Conditional probabilities are closer to the true structure than Transformer attention in all three tasks, and GCT matches or improves on GCNP.GCT improves upon the conditional-probability matrix for graph reconstruction and diagnosis-treatment classification; regularization toward that matrix helps on the other tasks.
- Attention Behavior Visualization: During graph reconstruction, GCT initially follows conditional probabilities and then shifts toward the true adjacency structure across self-attention blocks.The visualization shows data statistics guiding early attention while later blocks learn a distribution more similar to the true connections.
5 Conclusion
The paper proposes GCT to capture unknown encounter structure and reports stronger performance than baseline models on synthetic and publicly available EHR encounter tasks.
- GCT captures underlying encounter structure when the structure information is unknown.
- GCT outperformed various baseline models on encounter-based tasks using synthetic data and a publicly available EHR dataset.
- The authors identify GCT as a potential general-purpose EHR modeling algorithm.
A Generating Synthetic Encounter Records
The synthetic-record generation process models long-tailed medical codes and hierarchical diagnosis-treatment-lab relationships, while filtering encounters to maintain structural complexity and realism.
- Synthetic records could not be submitted as supplementary material because of their large volume.
- The records were generated by sampling diagnoses, likely co-occurring diagnoses, treatments, and lab codes according to specified distributions.
- Treatment codes were sampled conditionally on diagnoses, with treatments assumed conditionally independent given a diagnosis.
B Diagnosis-Treatment Classification Task
The diagnosis-treatment classification task evaluates whether a visit representation preserves encounter structure by predicting which diagnosis-treatment connections occur.
- The task is multi-label classification based on whether encounters contain specified connected diagnosis-treatment pairs.
- The task was made difficult because detecting labels from feature co-occurrence alone should not yield a perfect score.
- Sampling probabilities were adjusted so the two target connection types occurred with approximately equal probability.
C Training Details
Models were trained with common optimization settings, while architecture depths and hyperparameters were selected to support convergence and comparable parameterization.
- All models used Adam, minibatches of 32, and 1,000,000 iterations, with validation for model selection and test-set evaluation.
- The encounter-feature embedding size was set to 128 after preliminary experiments.
- Transformer used three self-attention blocks, while Shallow, Deep, GCN, GCNP, and GCNrandom depths were chosen to match trainable parameters.
- GCT used three self-attention blocks and one attention head, with learning rate, dropout, and regularization searched by Bayesian optimization.
D Prediction Performance in AUROC
Tables 7 and 8 report AUROC results for graph reconstruction, diagnosis-treatment classification, readmission prediction, and mortality prediction. GCT outperformed all other models across these tasks in AUROC.
- Table 7 reports graph reconstruction and diagnosis-treatment classification performance measured in AUROC.
- GCT outperformed all other models across graph reconstruction, diagnosis-treatment classification, readmission prediction, and mortality prediction in AUROC.Tables 7 and 8 report these task-specific comparisons.
- Table 8 reports readmission and mortality prediction performance on eICU measured in AUROC.
E Attention Behavior
The attention analyses compare Transformer and GCT on graph reconstruction and masked diagnosis code prediction. GCT begins with statistics-guided attention and progressively develops attention that can more closely reflect the true encounter structure, although this similarity is weaker for masked diagnosis prediction.
- Graph reconstruction: Transformer initially attends evenly across codes, recovers two true D 199 connections in its second block, then attends to incorrect codes in the third.
- Graph reconstruction: GCT begins graph reconstruction with attention shaped by scalar values and conditional probability matrix P, then progressively becomes more similar to the true adjacency matrix.The guided attention is especially evident for treatment codes such as T 939.
- Masked diagnosis code prediction: In masked diagnosis code prediction, GCT starts from P but develops attention maps less similar to the true structure than in graph reconstruction.
- Masked diagnosis code prediction: GCT attends heavily to the placeholder Visit node during masked diagnosis code prediction because diagnosis codes are allowed to attend only to treatment codes.
- Figure encoding: The figures encode diagnosis, treatment, and lab codes by prefixes D, T, and L; red bars show true connections, while blue bars show attention to all codes.Figures 5–8 focus on attention from D 199; Figures 7 and 8 additionally identify D 294 as the masked diagnosis code.