Book Review: Essential Graph RAG

Coming from a background of Knowledge Graph (KG) backed Medical Search, I don't need to be convinced about the importance of manually curated structured knowledge on the quality of search results. Traditional search is being rapidly replaced with Generative AI using a technique called Retrieval Augmented Generation (RAG), where the pipeline produces an answer summarizing the search results retrieved instead of the ten blue links that the searcher had to parse and retrieve an answer from earlier. In any case, I had been experimenting with Using KGs to enhance RAG to support this intuition, and when Microsoft announced their work on GtaphRAG, it felt good to be vindicated. So when Manning reached out to me to ask if I would be interested in reviewing the book Essential GraphRAG by Tomaž Bratanič and Oskar Hane, I jumped at the chance.

Both authors are from Neo4j, so it is not surprising that the search component is also Neo4j, even for vector search, and hybrid search is really vector + graph search (rather than the more common vector + lexical search). However, most people nowadays would prefer a multi-backend search that would include graph search as well as vector and lexical search, so the examples can help you learn (a) how to use Neo4j for vector search and (b) how to implement graph search with Neo4j. Since Neo4j is a leading graph database provider, this is useful information to know if you decide to incorporate graph search into your repertoire of tools, as you very likely are if you are reading this book.

The book is available under the Manning Early Access Program (MEAP) and is expected to be published in August 2025. It is currently organized into 8 chapters as follows:

Improving LLM accuracy -- here the authors introduce what LLMs are, what they are capable of as well as their limitations when used for question answering, i.e. not knowing about recent events post its training date, its tendency to hallucinate when it cannot answwe a question from the knowledge it was trained on, and its inability to know of company confidential or otherwise private information, since it is trained on public data only. They cover solutions to mitigate this, i.e. finetuning and RAG, and why RAG is a better alternaive in most cass. Finally they cover why KGs are the best general purpose datastore for RAG pipelines.

Vector Similarity Search and Hybrid Search -- here the authors cover the fundamentals of vector search, such as vector similarity functions, embedding models used to support vector search, and the reasoning behind chunking. They describe what a typical RAG pipeline looks like, although as mentioned earlier, they showcase Neo4j's vector search capabilities instead of relying on more popular vecror search alternatives. I thought it was good information though, since I wasn't aware that Neo4j supported vector search. They also cover hybrid search, in this case vector + graph search (this is a book about GraphRAG after all). Although I can definitely see Graph Search as one of the components of a hybrid search pipeline.

Advanced Vector Retrieval Strategies -- in this chapter, the authors introduce some interesting techniques to make your Graph Search produce more relevant context for your GraphRAG pipeline. Techniques on the query side include Step Back Prompting (SBP) to look for more generic concepts then drill down using Graph Search to improve recall, and the Parent Document Retriever pattern of retrieving parent documents of the chunks that matched, rather than the chunks themselves. On the indexing side, they talk about creating additional synthetic chunks that summarize actual chunks and can be queried as well as the chunks, and representing document chunks as pre-generated questions the chunk can answer instead of its text content.

Text2Cypher -- in this chapter, the authors show how an LLM can be prompted using Few Shot Learning (FSL) to generate Cypher queries from natural language. Users would type in a query using natural language, knowing nothing about the schema structure of the underlying Graph Database. The LLM, through detailed prompts and examples, would translate the natural language query to Cypher query. The authors also reference pre-trained models from Neo4j that have been fine-tuned to do this. While these models are generally not as effective as the one built from LLMs through prompting, they are more efficient on large volumes of data.

Agentic RAG -- Agentic RAG allows autonomous / semi-autonomous LLM backed software components, called Agents, to modify and enhance the standard control flow for RAG. One change could be for an Agent (the Router) to determine query intent and call on one or more retrieveers from the available pool of retrievers, or another (the Critic) to determine if the answer generated so far is adequate given the user's query, and if not, to rerun the pipeline with a modified query until the query is fully answered. The authors go on to describe a system (with code) consisting of a Router and Critic and several Retrieval Agents.

Constructing Knowledge Graph with LLM -- this chapter focuses on the index creation. Search is traditionally done on unstructured data such as text documents. This chapter describes using the LLM to extract entities of known types (PERSON, ORGANIZATION, LOCATION, etc), followed by a manual / semi-manual Graph Modeling step to set up relations between these extracted entities and build a schema. It then talks a little about convert specific query types into structured Cypher queries that leverage this schema.

Microsoft GraphRAG Implementation -- this chapter deals specifically with Microsoft's GraphRAG implementation. While most people think of GraphRAG as any infrastructure that supports incorporating Graph Search into a RAG pipeline, Microsoft specifies it as a multi-step recipe to build your KG from your data sources and use results from your KG to support a RAG pipeline. The steps involved are structured extraction and community detection, followed by summarization of community chunks into synthetic nodes. To some extent this is similar to Chonkie's Semantic Double Pass Merging (SDPM) chunker, except that the size of the skip window is unbounded. These synthetic chunks can be useful to answer global questions that span multiple ideas across the corpus. However, as the authors show, this approach can be effective for local queries as well.

RAG Application Evaluation -- because of the stochastic nature of LLMs, evaluating RAG pipelines in general present some unique challenges. Here these challenges are investigated with particular reference to GraphRAG systems, i.e. where the retrieval context is provided by Knowledge Graphs. The authors describe some metrics fro the RAGAS library, where LLMs are used to generate these metrics from outputs at different stages of the RAG pipeline. It also discusses ideas for setting up an evaluation dataset. The metrics covered in the example sare RAGAS context recall, faithfulness and answwr correctness.

Overall, the book takes a very practical, hands-on approach to the subject. It is filled with code examples and practical advice for leveraging KGs in RAG, and using Large Language Models (LLM) to build KGs, as well as evaluating such pipelines. If you were thinking of incorporating Graph Search into your search pipeline, be it traditional, hybrid, RAG or agentic, you will find the information in the book useful and beneficial.

KGC/HCLS 2024 Trip Report

I was at KGC (Knowledge Graph Conference) 2024, which is happening May 6-10 at Cornell Tech. I was presenting (virtually) at their Health Care and Life Sciences (HCLS) workshop, so my speakers pass was only valid for today for the HCLS portion of KGC. My trip report covers a few talks that I attended here. Attending virtually was a bit chaotic as sessions went over sometimes, so you might leave a session to attend another, only to find that it hadn’t started yet. This is hard to forsee, we have faced this issue ourselves the first time we moved an internal conference from in-person to hybrid.

KGs in RAG (Tom Smoker, WhatWhyHow.AI)

I have been working with Large Language Models (LLMs) and Retrieval Augmented Generation (RAG) for almost a year now, and I went to this talk hoping for insights on how to use graphs as input to RAG systems. Understandably, the speaker spent some time covering the basics, which I personally did not find very fruitful. However, there were some nuggets of wisdom I got out of the talk. First, the RAG pipelines can lower the risk of hallucinations by using LLMs for planning and reasoning, but without delegating to LLMs for factual information. And second, an agent architecture can more efficiently use smaller sub-graphs which can often be generated dynamically in Closed World models.

A side discussion on chat also yielded a paper reference Getting from Generative AI to Trustworthy AI: what LLMs may learn from Cyc (Lenat and Marcus, 2023). The paper looks really interesting on an initial skim and I plan to read in more detail later.

Knowledge Graphs for Precision Oncology (Krishna Bulusu, AstraZeneca)

A nice overview of applications of Knowledge Graph (KG) to Drug Discovery (DD). DD attempts to apply KG to solve 3 main problems: (1) find gene causing disease (2) match drug with disease and (3) (drug, gene, disease) as a fundamental relationship in DD. The speaker pointed out that the big advantage of KGs is Explainability. He also mentioned the use of graph clustering for node stratification.

Combining graph and vector representation for efficient information retrieval (Peio Popov, Ontotext)

This was a presentation from OntoText where they demonstrated new features built into their GraphDB database. This was of interest to me personally since our KG is also built using GraphDB. Specifically they have integrated LLM and vector search support into their products so they can be invoked from a SPARQL query. This gives GraphDB users the power to combine these techniques in the same call rather than build multi-stage pipelines.

I also learned the distinction between Semantic, Full text and Vector Search as ones based off KG, Lucene (or Lucene-like) indexes and vector search platforms, I would previously conflate the first and third.

Knowledge Engineering in Clinical Decision Support: When a Graph Representational Model is not enough (Maulik Kamdar, Optum)

This was a presentation from my ex-colleague Maulik Kamdar. He talks about challenges in Clinical Decision Support (CDS) where a KG alone is insufficient. Specifically the case he is considering where multiple third party ontologies need to be aligned into one KG. In this situation, similar concepts are combined into ValueSets, which are then composed with bare concepts or with each other to form Clinical Rules. Clinical Rules are further combined to form Clinical Calculators or Questionnaires, which are then combined to form Decision Trees and Flowcharts, which are then combined into Clinical Guidelines. I am probably biased given our common history, but I found this talk to be the most educational for me.

Knowledge Graphs, Theorem Provers and Language Models (Vijay Saraswat and Nikolaos Vasiloglou)

The speakers discussed the role of self-discovery, In-Context Learning (ICL), symbiotic integration of KG with search, and Graph RAG in reasoning engines powered by KG and LLM. They characterize an Agent as an LLM based black box that is provided with pairs of input-output instances to learn some unknown function (similar to ML models). They describe ICL as learning through few shot and many shot examples. They also talk about using the output of KG to fact-check / enhance LLMs and using LLMs to generate assertions that can be used to create a KG. Their demo shows how an LLM is able to learn to generate a Datalog like graph query language from text prompts using few-shot examples.

The speaker made reference to the following three papers in support of the techniques he was describing, which I have duly added to my reading list.

• From words to numbers: Your Large Language Model is secretly a capable regressor when given in-context examples (Vacareanu et al, 2024)
• Many-shot In-Context Learning (Agarwal et al, 2024)
• Self-Discover: Large Language Models self-compose Reasoning Structures (Zhou et al, 2024)

A Scalable and Robust Named Entity Recognition and Linking System for a Clinical Healthcare Knowledge Graph (Sujit Pal, Elsevier Health)

This was my talk. I had originally intended to attend in person but it seemed wasteful to fly across the country to deliver a 5-minute presentation. It did take a bit of planning to present remotely but I learned two useful life lessons.

1. You can generate a presentation video from MS Powerpoint. Simply create your slides and record a slideshow where you record yourself narrating your presentation. Once done, export as an MP4 and upload to Youtube or other video service.
2. You can print posters online and have them delivered to someone else.

Huge thanks to my colleague Tom Woodcock who attended in person, and who was kind enough to carry and hang my poster at the conference for me, and who also agreed to present my slideshow for me (although I think that in the end he did not have to). Many thanks also to my ex-colleague Helena Deus (part of the HCLS organizing team), who helped walk me through to a workable solution and was instrumental in my talk being delivered successfully. Also thanks to Leah Walton from the HCLS organizing team, for supporting me in my attempt to present remotely.

Here is the Youtube video for my 5-minute presentation in case you are interested. It’s a bit high-level since I had only 5 minutes to cover everything, but there is a little more information in the poster below.

Graphs for good – Hypothesis generation for Rare Disease Treatment (Brian Martin, AbbVie)

This presentation revolves around a graph that connects diseases to drugs via disease variants, gene, pathway, gene and compound entities. This was used to find a cure for a rare disease using existing medications. It was later extended to find candidate cures for a group of 20 most neglected diseases worldwide. The speakers verified that results for Dengue fever correlates well with previously known information, thus supporting the veracity of the approach. The paper describing this work is Leveraging a Billion-Edge Knowledge Graph for Drug Re-purposing and Target Prioritization using Genomically-Informed Subgraphs (Martin et al, 2022).

Generating and Querying Graphs with LLM (Brian Martin, Subha Madhavan, Berenice Wulbrecht)

Panel discussion where various strategies for generating and querying graphs using LLMs were discussed. Entertaining (and somewhat predictable) comparisons of Property Graphs vs RDF graphs to Ford and Ferrari automobiles, and how LLMs transform them into Teslas (with its self-driving technology). They also talk about extracting assertions from a corpus of documents to create a KG customized for the corpus, and then using the KG to fact-check the output of the LLM for RAG queries against that corpus.

Overall, I think it was a great conference. Learned a lot, would love to go back and present here in the future, hopefully this time in person.

Knowledge Graph Aligned Entity Linker using SentenceTransformers

Most of us are familiar with Named Entity Recognizers (NERs) that can recognize spans in text as belonging to a small number of classes, such as Person (PER), Organization (ORG), Location (LOC), etc. These are usually multi-class classifier models, trained on input sequences to return BIO (Begin-Input-Output) tags for each token. However, recognizing entities in a Knowledge Graph (KG) using this approach is usually a much harder proposition, since a KG can contain thousands, even millions, of distinct entities, and it is just not practical to create a multi-class classifier for so many target classes. A common approach to building a NER for such a large number of entities is to use dictionary based matching. However, the approach suffers from the inability to do "fuzzy" or inexact matching, beyond standard normalization streategies such as lowercasing and stemming / lemmatizing, and requires you to specify up-front all possible synonyms that may be used to refer to a given entity.

An alternative approach may be to train another model, called a Named Entity Linker (NEL) that would take the spans recognized as candidate entities or phrases by the NER model, and then attempt to link the phrase to an entity in the KG. In this situation, the NER just learns to predict candidate phrases that may be entities of interest, which puts it on par with simpler PER/ORG/LOC style NERs in terms of complexity. The NER and NEL are pipelined together in a setup that is usually known as Named Entity Recognition and Linking (NERL).

In this post, I will describe a NEL model that I built for my 2023 Dev10 project. Our Dev10 program allows employees to use up to 10 working days per year to pursue a side-project, similar to Google's 20% program. The objective is to learn an embedding model where encodings of synonyms of a given entity are close together, and where encodings of synonyms of different entities are pushed far apart. We can then encode each entity in this space as the encoding of the centroid of the encodings of its individual synonyms. Each candidate phrase output from the NER model can then be encoded using this embedding model, and its nearest neighbors in the embedding space would correspond to the most likely entities to link to.

The idea is inspired by Self-Alignment Pretraining for Biomedical Entity Representations (Liu et al, 2021) which produced the SapBERT model (SAP == Self Aligned Pretraining). It uses Contrastive Learning to fine-tune the BiomedBERT model. In this scenario, positive pairs are pairs of synonyms for the same entity in the KG and negative pairs are synonyms from different entities. It uses the Unified Medical Language System (UMLS) as its KG, to source synonym pairs.

I follow a largely similar approach in my project, except that I use the SentenceTransformers library to fine tune the BiomedBERT model. For my initial experiments, I also used the UMLS as my source of synonym pairs, mainly for reproducibility purposes since it is a free resource available for download to anyone. I tried fine-tuning a bert-base-uncased model and the BiomedBERT models, with MultipleNegativesRanking (MNR) as well as Triplet loss, the latter with Hard Negative Mining. My findings are in line with the SapBERT paper, i.e. that BiomedBERT performs better than BERT base, and that MNR performs better than Triplet loss. The last bit was something of a dissapointment, since I had expected Triplet loss to perform better. It is possible that the Hard Negative Mining was not hard enough, or maybe I needed a higher number than 5 negatives for each positive.

You can learn more about the project in my GitHub repository sujitpal/kg-aligned-entity-linker, as well as find the code in there, in case you want to replicate it.

Here are some visualizations from my best model. The chart on the left shows the distribution of cosine similarities between known negative synonym pairs (orange curve) and known positive synonym pairs (blue curve). As you can see, there is almost no overlap. The heatmap on the right shows the cosine similarity of a set of 10 synonym pairs, where the diagonal corresponds to positive pairs and the non-diagonal elements correspond to negative pairs. As you can see, the distribution seems quite good.

I also built a small demo that shows what in my opinion is the main use case for this model. It is a NERL pipeline, where the NER component is the UMLS entity finder (en_core_sci_sm) from the SciSpacy project, and the NEL component is my best performing model (kgnel-bmbert-mnr). In order to look up nearest neighbors for a given phrase encoding, the NEL component also needs a vector store to store the centroids of the encodings of entity synonyms, I used QDrant for this purpose. The QDrant vector store needs to be populated with the centroid embeddings in advance, and in order to cut down on the index and vectorization time, I only computed embeddings for centroids for entities of type "Disease or Syndrome" and "Clinical Drug". The visualizations below show the outputs (from displacy) of the outputs of the NER component:

and that of the NEL component in my demo NERL pipeline. Note that only spans that were identified as a Disease or Drug with a confidence above a threshold were selected in this phase.

Such a NERL pipeline could be used to mine new literature for new synonyms of existing entities. Once discovered, they could be added to the synonym list for the dictionary based NER to increase its recall.

Anyway, that was all I had for this post. Today is also January 1 2024, so I wanted to wish you all a very Happy New Year and a productive 2024 filled with many Machine Learning adventures!