Source-linked AI summary

A Method to Learn Embedding of a Probabilistic Medical Knowledge Graph: Algorithm Development

Linfeng Li, Peng Wang, Yao Wang, Jinpeng Jiang, Buzhou Tang, Jun Yan, Shenghui Wang, Yuting Liu

arXiv:1909.00672v2cs.AIcs.CL

TL;DR

Medical EMR knowledge graphs contain triplets with graded uncertainty, challenging embedding methods built for binary-correct deterministic facts. The paper proposes PrTransX to incorporate triplet probabilities into embeddings and reports better link-prediction performance than corresponding TransX models, while noting that application in clinical decision support remains future work.

  • Problem

    Medical EMR triplets have probabilistic certainty rather than binary correctness, but existing TransX algorithms do not incorporate triplet probabilities.

    Source: Related Approaches of Knowledge Graph… · p. 3
  • Method

    PrTransX extends TransX by mapping scores to triplet probabilities and training embeddings with probability-based loss on probabilistic medical knowledge graphs.

    Source: Objective · p. 3
  • Results

    PrTransX significantly outperformed corresponding TransX algorithms on link prediction in a medical knowledge graph built from large-scale EMR data.

    Source: Conclusion · p. 11
  • Takeaways & Limitations

    The results support incorporating uncertainty into embeddings for probabilistic medical knowledge graphs.

    Source: Conclusion · p. 11

    Applying PrTransX embeddings in clinical decision support remains future work, and raw EMR test tails can be noisy without expert ground-truth labeling.

    Source: Data Set Split and Ground Truth Setup · p. 5, Future Studies · p. 11

Abstract

from arXiv · show

This paper proposes an algorithm named as PrTransH to learn embedding vectors from real world EMR data based medical knowledge. The unique challenge in embedding medical knowledge graph from real world EMR data is that the uncertainty of knowledge triplets blurs the border between "correct triplet" and "wrong triplet", changing the fundamental assumption of many existing algorithms. To address the challenge, some enhancements are made to existing TransH algorithm, including: 1) involve probability of medical knowledge triplet into training objective; 2) replace the margin-based ranking loss with unified loss calculation considering both valid and corrupted triplets; 3) augment training data set with medical background knowledge. Verifications on real world EMR data based medical knowledge graph prove that PrTransH outperforms TransH in link prediction task. To the best of our survey, this paper is the first one to learn and verify knowledge embedding on probabilistic knowledge graphs.

Introduction

Medical knowledge graphs support clinical decision systems, but EMR-derived triplets are probabilistic rather than absolutely correct or wrong. Existing embedding approaches model deterministic facts, motivating probabilistic extensions of translation-based methods.

Background

  • Medical knowledge graphs support clinical decision support, but expert construction is labor-intensive, making EMR mining a promising scalable alternative.Source: Background · p. 2

Related Approaches of Knowledge Graph Embedding

Objective

The study proposes PrTransX, an extension of classical TransX algorithms for learning embeddings from probabilistic knowledge graphs.

Objective

  • PrTransX extends classical TransX algorithms to learn embeddings for probabilistic knowledge graphs.Source: Objective · p. 3

Mapping Function Between Score Value and Probability

The method connects TransX score values with triplet probabilities, trains embeddings to reproduce observed probabilities, and constructs probabilities from EMR visit records. It applies this framework to large-scale medical data.

Mapping Function Between Score Value and Probability

Loss Function to Learn Embedding of a Probabilistic Knowledge Graph

Evaluation of the Proposed Algorithms on Large-Scale Real-World Electronic Medical Records Data

Building the Knowledge Graph From Real-World Electronic Medical Records

Data Set Split and Ground Truth Setup

The data split keeps all tails for a given head–relation pair together, while expert labeling supplies a clinically grounded evaluation subset because raw EMR tails can be noisy.

Data Set Split and Ground Truth Setup

Negative Triplets Sampling

Negative sampling uses entity types and medical background knowledge to reduce false negatives, while a mapping function converts TransX scores into probabilistic triplet values.

Negative Triplets Sampling

  • Bernoulli sampling determines whether to replace the head or tail, lowering false negatives in TransH-based algorithms.Source: Negative Triplets Sampling · p. 5
  • Entity-type constraints and medical knowledge reduce false-negative triplets compared with purely statistical sampling.Candidate replacements are restricted to valid relation-specific entity types, and medically related entities such as ICD-10 hyponyms are excluded from corruption.Source: Negative Triplets Sampling · p. 5

Mapping Function Between the Score Value and Probability

  • Medical triplet labels represent probabilities rather than absolute correctness, ranging from 0 when a triplet never holds to 1 when it always holds.Source: Mapping Function Between the Score… · p. 5
  • The mapping function Φ() links TransX distance scores to probabilities, enabling probability inference from embedding vectors alone and independent triplet probabilities.The construction supports unseen-triplet probability estimation and avoids softmax limitations for 1-to-many and many-to-1 relations.Source: Mapping Function Between the Score… · pp. 5–6

Loss Function to Learn Embedding of the Probabilistic Knowledge Graph

The proposed loss extends margin-based training by matching scores to observed triplet probabilities for both valid and negative triplets, with link prediction evaluated using ranking and relevance-sensitive metrics.

Loss Function to Learn Embedding of the Probabilistic Knowledge Graph

  • PrTransX combines margin-based loss with probability-based losses for positive and negative triplets.The probability terms require scores to approximate the mapped statistical probabilities, addressing the margin loss's failure to calibrate predicted probabilities.Source: Loss Function to Learn Embedding of… · p. 6

Evaluation Protocol

  • Link prediction ranks valid tail entities for given heads and relations, using Hits@10, mean rank, and NDCG@10 for evaluation.NDCG@10 accounts for differing relevance levels among medical tail entities.Source: Evaluation Protocol · p. 7

Real-World Electronic Medical Records Data and Knowledge Graph

The study builds a medical knowledge graph from large-scale hospital EMR data, incorporates ICD-10 hierarchy and prior medical relations, and selects high-probability triplets for training and evaluation.

Real-World Electronic Medical Records Data and Knowledge Graph

Data Set

  • The dataset contains 3,767,198 patients and 16,217,270 visits collected from 2015 to 2018, with disease entities identified by ICD-10 codes.ICD-10 hierarchy generates triplets at multiple disease-code levels.Source: Data Set · p. 7
  • The graph includes five extracted relationships plus upper_disease_to_lower_disease as prior domain knowledge, while retaining only the top 20 triplets by probability.Source: Data Set · p. 7
  • The separated data contain 205,877 training triplets from 21,327 groups and 49,756 test triplets from 4,547 groups.Source: Data Set · p. 7
  • The evaluation set contains 335 triplets from 25 groups distributed evenly across five relationships.Source: Data Set · p. 7

Baselines and Implementation

The experiments compare TransX models with their probabilistic extensions under shared embedding and optimization settings, while naive Bayes cannot support the group-based data split for tail prediction.

Baselines and Implementation

Evaluation Results

PrTransX variants were compared with corresponding TransX models across overall and relation-specific link-prediction evaluations. The probabilistic variants generally achieved better aggregate performance, while relation statistics showed substantial differences in entity connectivity.

Overall comparison

  • PrTransX outperformed corresponding TransX models for X=E/H/R/D/Sparse in nearly all indicators, with higher Hits@10 and NDCG@10 and lower mean rank.Scores were averaged over all triplets in the evaluation set.Source: Evaluation Results · p. 8
  • The evaluation compared TransX and PrTransX using Hits@10, mean rank, and NDCG@10 across all evaluation-set triplets.Source: Evaluation Results · p. 8, Figure 3 · p. 8, Figure 4 · p. 8

Relation-specific analysis

Relation distributions

  • The relation disease_to_examination had a median near 150 head entities per tail and a 75th percentile above 350.This means 50% of examination entities were related to near 150 diseases, while 25% were related to more than 350 diseases.Source: Evaluation Results · p. 9, Figure 6 · p. 9
  • Disease-to-examination was less specific than other relations, whereas disease_to_operation was the most specific with a median of 10.The observed pattern matches the stated medical interpretation that examinations apply to many diseases, while operations apply to fewer.Source: Evaluation Results · p. 9

Discussion

PrTransX generally improved link prediction over corresponding TransX models, while relation structure and model-specific geometry shaped performance differences. Probability-based training was especially beneficial for PrTransH/R/D, whereas PrTransE showed little change.

Principal Findings

  • PrTransX outperformed corresponding TransX models across nearly all indicators, with higher Hits@10 and NDCG@10 and lower mean rank.PrTransR performed best under Hits@10 and NDCG@10, while PrTransE performed best under mean rank.Source: Principal Findings · pp. 10–11
  • The medical graph’s five relations favored TransE over TransH/R/D because all entities and relations could be trained in one shared space.The authors relate this pattern to TransH’s weaker performance than TransE on WN18 [6].Source: Principal Findings · p. 10
  • PrTransE and TransE had similar Hits@10 and NDCG@10, whereas probability-based targets significantly improved PrTransH/R/D over their corresponding baselines.The shared-space constraint made the probability targets of different relations difficult to satisfy for TransE.Source: Principal Findings · p. 10
  • Relation specificity affected difficulty: disease_to_examination was broad and difficult to predict, whereas disease_to_operation was specific and easier.The same examination can apply to many diseases, while the same operation applies to relatively few diseases.Source: Principal Findings · p. 11
  • PrTransX outperformed TransX even for the difficult disease_to_examination relation, while the largest gains included PrTransR on disease_to_operation.PrTransD/Sparse(share)/Sparse(separate) also significantly improved over their corresponding TransX models for disease_to_operation.Source: Principal Findings · p. 11

Future Studies

The authors identify broader applications for probabilistic knowledge embeddings beyond link prediction and conclude that PrTransX successfully incorporates triplet uncertainty. Applying these embeddings in clinical decision support remains future work.

Future Studies

  • PrTransX embeddings could be applied beyond link prediction, including clinical decision support, but this application remains to be explored.Prior work used TransE and latent-factor embeddings to infer diagnoses from laboratory results and symptoms [11].Source: Future Studies · p. 11

Conclusion

  • The study concludes that PrTransX significantly outperformed corresponding TransX algorithms on link prediction using a medical KG constructed from large-scale EMR data.The authors interpret this as successful incorporation of knowledge-triplet uncertainty into embedding vectors.Source: Conclusion · p. 11

Abbreviations

  • The paper defines EMR as electronic medical record, KG as knowledge graph, and PBE as probabilistic belief embedding.It also defines HITS@10 and NDCG@10 as evaluation measures.Source: Abbreviations · p. 12
  • The article was published in JMIR Medical Informatics in 2020 under a Creative Commons Attribution License.The supplied publication record identifies the article as JMIR Med Inform 2020;8(5):e17645.Source: Abbreviations · p. 12
Loading 1909.00672v2…