Benchmark and Best Practices for Biomedical Knowledge Graph Embeddings (2020)

Abstract: Much of biomedical and healthcare data is encoded in discrete, symbolic form such as text and medical codes. There is a wealth of expert-curated biomedical domain knowledge stored in knowledge bases and ontologies, but the lack of reliable methods for learning knowledge representation has limited their usefulness in machine learning applications. While text-based representation learning has significantly improved in recent years through advances in natural language processing, attempts to learn biomedical concept embeddings so far have been lacking. A recent family of models called knowledge graph embeddings have shown promising results on general domain knowledge graphs, and we explore their capabilities in the biomedical domain. We train several state-of-the-art knowledge graph embedding models on the SNOMED-CT knowledge graph, provide a benchmark with comparison to existing methods and in-depth discussion on best practices, and make a case for the importance of leveraging the multi-relational nature of knowledge graphs for learning biomedical knowledge representation. The embeddings, code, and materials will be made available to the community.

## introduction a vast amount of biomedical domain knowledge is stored in knowledge bases and ontologies. for example, snomed clinical terms (snomed-ct) foot_1 is the most widely used clinical terminology in the world for documentation and reporting in healthcare, containing hundreds of thousands of medical terms and their relations, organized in a polyhierarchical structure. snomed-ct can be thought of as a knowledge graph: a collection of triples consisting of a head entity, a relation, and a tail entity, denoted (h, r, t). snomed-ct is one of over a hundred terminologies under the unified medical language system (umls) @xcite , which provides a metathesaurus that combines millions of biomedical concepts and relations under a common ontological framework. the unique identifiers assigned to the concepts as well as the resource release format (rrf) standard enable interoperability and reliable access to information. the umls and the terminologies it encompasses are a crucial resource for biomedical and healthcare research. one of the main obstacles in clinical and biomedical natural language processing (nlp) is the ability to effectively represent and incorporate domain knowledge. a wide range of downstream applications such as entity linking, summarization, patientlevel modeling, and knowledge-grounded language models could all benefit from improvements in our ability to represent domain knowledge. while recent advances in nlp have dramatically improved textual representation @xcite , attempts to learn analogous dense vector representations for biomedical concepts in a terminology or knowledge graph (concept embeddings) so far have several drawbacks that limit their usability and wide-spread adoption. further, there is currently no established best practice or benchmark for training and comparing such embeddings. in this paper, we explore knowledge graph embedding (kge) models as alternatives to existing methods and make the following contributions: ## visualization we first discuss the embedding visualizations obtained through largevis @xcite , an efficient large-scale dimensionality reduction technique available as an application in graphvite. figure 1 shows concept embeddings for rotate, complex, snomed2vec, and cui2vec, with colors corresponding to broad semantic groups. cui2vec embeddings show structure but not coherent semantic clusters. snomed2vec shows tighter groupings of entities, though the clusters are patchy and scattered across the embedding space. complex produces globular clusters centered around the origin, with clearer boundaries between groups. rotate gives visibly distinct clusters with clear group separation that appear intuitive: entities of the physiology semantic group (black) overlap heavily with those of disorders (magenta); also entities under the concepts semantic group (red) are relatively scattered, perhaps due to their abstract nature, compared to more concrete entities like devices (cyan), anatomy (blue), and chemicals (green), which form tighter clusters. interestingly, the embedding visualizations for the 5 kge models fall into 2 types: rotate and transe produce well-separated clusters while com-plex, distmult and simple produce globular clusters around the origin. since the plots for each type appear almost indistinguishable we show one from each (rotate and complex). we attribute the characteristic difference between the two model types to the nature of their scoring functions: ro-tate and transe have an additive component while complex, distmult and simple are multiplicative. figure 2 shows more fine-grained semantic structure by coloring 5 selected semantic types under the procedures semantic group and greying out the rest. we see that rotate produces subclusters that are also intuitive. laboratory procedures are well-separated on their own, health care activity and educational activity overlap significantly, and diagnostic procedures and therapeutic or preventative procedures overlap significantly. complex also reveals subclusters with globular shape, and snomed2vec captures laboratory procedures well but leaves other types scattered. these observations are consistent across other semantic groups. we include similar visualizations for the chemicals & drugs semantic group in the supplements. while semantic class information is not the only significant aspect of snomed-ct, since the snomed-ct graph is largely organized around semantic group and type information, it is promising that embeddings learned (without supervision) preserve it. no previous results to compare to, we include performance of transe and rotate on two standard kge benchmark datasets for reference: fb15k-237 (14,541 entities, 237 relations, and 310,116 triples) and wn18rr (40,943 entities, 11 relations, and 93,003 triples). given that snomed-ct is larger and arguably a more complex knowledge graph than the two datasets, the link prediction results suggest that the kge models learn a reasonable representation of snomed-ct. we include sample model outputs for the top 10 entity scores for link prediction in the supplements. ## discussion based on our findings, we recommend the use of kge models to leverage the multi-relational nature of knowledge graphs for learning biomedical concept and relation embeddings; and of appropriate evaluation tasks such as link prediction, entity classification and relation prediction for fair comparison across models. we also encourage analysis beyond standard validation metrics, e.g. visualization, examining model predictions, reporting metrics for different relation groupings and devising problem or domain-specific validation tasks. a further promising evaluation task is the triple prediction proposed in @xcite , which we leave for future work. a more ideal way to assess concept embeddings in biomedical nlp applications and patient-level modeling would be to design a suite of benchmark downstream tasks that incorporate the embeddings, but that warrants a rigorous paper of its own and is left for future work. we believe this paper serves the biomedical nlp community as an introduction to kges and their evaluation and analyses, and also the kge community by providing a potential standard benchmark dataset with real-world relevance. ## conclusion and future work we present results from applying 5 leading kge models to the snomed-ct knowledge graph and compare them to related work through visualizations and evaluation tasks, making a case for the importance of using models that leverage the multirelation nature of knowledge graphs for learning biomedical knowledge representation. we discuss best practices for working with biomedical knowledge graphs and evaluating the embeddings learned from them, proposing link prediction, entity classification, and relation prediction as standard evalua-tion tasks. we encourage researchers to engage in further validation through visualizations, error analyses based on model predictions, examining stratified metrics, and devising domain-specific tasks that can assess the usefulness of the embeddings for a given application domain. there are several immediate avenues of future work. while we focus on the snomed-ct dataset and the kge models implemented in graphvite, other biomedical terminologies such as the gene ontology (the gene ontology consortium, 2018) and rxnorm @xcite could be explored and more recent kge models, e.g. tucker @xcite and murp @xcite , applied. additional sources of information could also potentially be incorporated, such as textual descriptions of entities and relations. in preliminary experiments, we initialized entity and relation embeddings with the embeddings of their textual descriptors extracted using clinical bert @xcite , but it did not yield gains. this may suggest that the concept and language spaces are substantially different and strategies to jointly train with linguistic and knowledge graph information require further study. other sources of information include entity types (e.g. umls semantic type) and paths, or multi-hop generalizations of the 1-hop relations (triples) typically used in kge models @xcite . notably, coke trains contextual knowledge graph embeddings using path-level information under an adapted version of the bert training paradigm @xcite . lastly, the usefulness of biomedical knowledge graph embeddings should be investigated in downstream applications in biomedical nlp such as information extraction, concept normalization and entity linking, computational fact checking, question answering, summarization, and patient trajectory modeling. in particular, entity linkers act as a bottleneck between text and concept spaces, and leveraging kges could help develop sophisticated tools to parse existing biomedical and clinical text datasets for concept-level annotations and additional insights. well performing entity linkers may then enable training knowledge-grounded largescale language models like knowbert @xcite . overall, methods for learning and incorporating domain-specific knowledge representation are still at an early stage and further discussions are needed.