Source-linked AI summary
Pre-training of Graph Augmented Transformers for Medication Recommendation
Junyuan Shang, Tengfei Ma, Cao Xiao, Jimeng Sun
TL;DR
Medication recommendation research is limited by selection bias from underusing single-visit EHR data and by insufficient incorporation of hierarchical medical-code knowledge. G-BERT combines graph neural networks with a BERT-style Transformer encoder, pre-trains on single-visit EHR data, and is then fine-tuned for medication recommendation, where it outperforms all baselines in prediction accuracy.
Problem
Existing medication recommendation work often uses longitudinal EHRs from patients with multiple visits, excluding many single-visit patients, and rarely embeds medical-code hierarchies in their original graph form.
Method
G-BERT learns medical-ontology embeddings with graph neural networks, integrates them into a Transformer-based visit encoder, and applies modified BERT-style self-prediction and dual-prediction pre-training.
Results
G-BERT outperforms all baselines in prediction accuracy on the medication recommendation task.
Takeaways & Limitations
G-BERT demonstrates that language-model pre-training and medical-ontology information can be combined for medical-code representation and medication recommendation.
Takeaways & Limitations
Future work includes adapting G-BERT to larger datasets with more heterogeneous modalities.
Abstract
from arXiv · showhide
Medication recommendation is an important healthcare application. It is commonly formulated as a temporal prediction task. Hence, most existing works only utilize longitudinal electronic health records (EHRs) from a small number of patients with multiple visits ignoring a large number of patients with a single visit (selection bias). Moreover, important hierarchical knowledge such as diagnosis hierarchy is not leveraged in the representation learning process. To address these challenges, we propose G-BERT, a new model to combine the power of Graph Neural Networks (GNNs) and BERT (Bidirectional Encoder Representations from Transformers) for medical code representation and medication recommendation. We use GNNs to represent the internal hierarchical structures of medical codes. Then we integrate the GNN representation into a transformer-based visit encoder and pre-train it on EHR data from patients only with a single visit. The pre-trained visit encoder and representation are then fine-tuned for downstream predictive tasks on longitudinal EHRs from patients with multiple visits. G-BERT is the first to bring the language model pre-training schema into the healthcare domain and it achieved state-of-the-art performance on the medication recommendation task.
1 Introduction
Existing medication recommendation models face selection bias from excluding single-visit patients and underuse hierarchical medical-code knowledge. G-BERT addresses both limitations by combining BERT-style pre-training with graph neural networks for medical-code representation.
- Medication recommendation models learn medical-entity representations from EHR data to predict medications suited to patients’ health conditions.
- Existing work has two limitations: selection bias from discarding patients with only one hospital visit and limited use of medical-code hierarchies.
- G-BERT combines pre-training techniques and graph neural networks to improve medical-code representation and medication recommendation.
- Its BERT-based pre-training operates on individual EHR visits, allowing single-visit data to be used.
- G-BERT learns medical-ontology embeddings with graph neural networks and feeds them into a multi-layer Transformer for BERT-style pre-training and fine-tuning.
2 Related Work
Prior medication recommendation work includes instance-based and longitudinal approaches, while related research also covers pre-training and graph neural networks. These strands motivate G-BERT’s combination of EHR pre-training and ontology-aware representation learning.
- Medication Recommendation: Medication recommendation methods are categorized into instance-based methods focused on current conditions and longitudinal methods using temporal dependencies.
- Pre-training Techniques: Pre-training techniques provide model-training initializations and unsupervised pre-training can support better generalization from training data.
- Graph Neural Networks (GNN): Graph neural networks learn node or graph representations from graph-structured data, including graph convolutional, message-passing, and graph-attention networks.
- Graph Neural Networks (GNN): GRAM represents medical concepts through their ontology ancestors with attention, whereas G-BERT differs from it in two aspects described in the paper.
3 Problem Formalization
The problem formalization represents each patient as a sequence of multivariate visits and each visit through diagnosis codes. Given a patient’s history, the task is to recommend multiple medications as a multi-label output.
- Each patient is represented as a sequence of multivariate observations, with T^(n) denoting the number of visits for patient n.
- Each visit is represented using its corresponding diagnosis-code set, with medical codes unified across code types.
- Medical codes are organized in tree-structured systems such as ICD-9 for diagnosis and ATC for medication, denoted by separate ontologies.
- Given the visit at time t and patient history X_1:t, the model recommends multiple medications as a multi-label output ŷ_t ∈ {0,1}^{|C^m|}.
4 Method
G-BERT combines graph-based ontology embedding with a shared-weight Transformer encoder and downstream classifier. It uses BERT-style pre-training objectives for medical codes before fine-tuning medication recommendation.
- Visit encoder: The framework feeds separate diagnosis and medication ontology embeddings into a shared-weight, multi-layer Transformer to produce visit representations.A [CLS] token provides the visit representation, while the Transformer aggregates unordered medical codes within each visit.
- Fine-tuning: For medication recommendation, G-BERT concatenates aggregated previous-visit embeddings with the current visit representation before fine-tuning prediction layers.The framework aggregates visit embeddings from diagnoses and medications separately before combining them with the current diagnosis visit embedding.
- Ontology embedding: G-BERT first derives medical-code ontology embeddings with graph neural networks, incorporating information from diagnosis and medication hierarchies.The ontology embeddings are computed in two stages, passing information from children to ancestors and then from ancestors back to leaf codes.
- Ontology embedding: The ontology graph updates ancestor nodes from their children and leaf nodes from their connected ancestors, unlike methods that update only leaf nodes.G-BERT initializes node embeddings with a learnable matrix and uses a two-step update procedure for both leaf and ancestor nodes.
- Pre-training: During pre-training, masked medical codes are restored using self-prediction within the same code type and dual-prediction across diagnosis and medication types.The dual-prediction objective adapts BERT’s next-sentence idea to predicting unknown medical codes from the other code type.
5 Experiment
The experiments evaluate G-BERT on MIMIC-III medication recommendation, comparing it with baseline models and ablated variants. Results indicate that combining ontology information with pre-training improves performance relative to ablations and baselines.
- Experimental Setting: Experiments use MIMIC-III data, with single-visit and multi-visit records used for pre-training and multi-visit patients used for evaluation.Multi-visit records are split into visit slices, and duplicate codes within visits are removed to avoid leakage.
- Baselines and Ablations: G-BERT is compared with logistic regression, LEAP, RETAIN, GRAM, GAMENet, and three ablated variants.The variants separately remove ontology information, pre-training, or both.
- Metrics: Medication recommendation accuracy is measured using Jaccard Similarity, Average F1, and Precision Recall AUC.Jaccard measures the intersection-over-union of the ground-truth and predicted medication sets.
- Ablation Results: G-BERTG−,P− performs worse than G-BERTG− and G-BERTP−, supporting the effectiveness of ontology-enhanced embeddings and unsupervised pre-training.The comparison is reported across the medication recommendation results.
- Baseline Results: G-BERT outperforms RETAIN and GAMENet, while GAMENet remains worse even with additional DDI knowledge and procedure codes.G-BERTP− also scores slightly higher than GRAM across all reported metrics.
6 Conclusion
The paper concludes that G-BERT combines language-model pre-training with graph-based medical ontology information for medication recommendation. It uses otherwise discarded single-visit EHR data and identifies broader auxiliary tasks and heterogeneous datasets as future directions.
- Conclusion: G-BERT adapts BERT to EHR data and integrates medical ontology information using graph neural networks for medical code representation and medication recommendation.The model additionally pre-trains on EHR from patients with only one hospital visit.
- Conclusion: G-BERT outperforms all baselines in medication recommendation prediction accuracy after pre-training on single-visit EHR data.The conclusion describes single-visit records as generally discarded before model training.
- Future Work: Future work includes adding auxiliary and structural tasks and adapting the model to larger datasets with more heterogeneous modalities.These directions are presented as ways to improve code representation and broaden the model’s data suitability.