Improving RAG with Knowledge Graph and Milvus

Bay.Area.AI: Improving RAG with Knowledge Graph and Milvus, Jiang Chen, Zilliz

Recording: Bay.Area.AI: Improving RAG with Knowledge Graph and Milvus, Jiang Chen, Zilliz

uh thanks for the introduction Alexi and thanks for inviting me here um so I'm Jang um I'm um a bit of introduction of myself uh I'm uh leading the ecosystem Integrations and developer relation as zist uh for uh folks that don't know like we are actually the creator of the open source Vector database called muas I know we have two names so that it's always confusing U but we are the company who are um uh still the the large uh the the major maintainer of the M project the M project itself has been donated to Linux Foundation data and AI um so the QR code is uh my LinkedIn if you want to connect with me in person we can uh we can talk more there um so the today I want to share more about the advanced m in rag um as um as most people have known like um rag has been really widely adopted in those AI applications however people are constantly frustrated by the quality problem of rag because naive rag um it is impressive but it doesn't solve all of the problem um so actually I came from the search indexing and infrastructure background um back in the days when I was building uh image and video search at Google um we actually uh heavily depend on two pieces of Technology uh to improve the ritual quality of of the search system one is the um uh deep neural network and the vector edings uh which has been known by many people already the other one is actually kg uh with uh Knowledge Graph so um there is a lot of things in knowledge engineering uh regarding like how do you systematically um extract the the knowledge information from the instructor data from the say WEP data corpers and then um organize them in a systematic way so that during serving time like online serving time you can effectively um retrieve and leverage the information behind that because you know for um just look at the how um like human beings are encoding the information right we we have the human language which is a great tool but the human language um usually talks about things together rather than giving it a structured uh representation so that having Vector embedding can help us retrieval uh to retrieve the RW text from the say um web pages your internal documentations um but it doesn't really uh solve all of the problem because um sometimes we still need some some structured way to encode the information so that you understand um you know we can help machine to understand the uh relationship between the knowledge entities the relationship between the concepts so that's a a great enhancement to to Rack to see this difference in action um so here is a a graph from from the paper which does the survey of the um you know Alm based retri Alm based uh uh chat and rack based and also kg plus rack based so as you can see that um you know simply using you're pretty much relying on the sole kind of essential capability of the large D gr model itself which was not trained on your data um so that adding I I'm with the retrieval or like vector search um you can leverage your Internal Documentation like things not a like information not available to the uh to the public or to the public data size where the large language model was trained on however um with only this you're pretty much relying on the combination of information um applied in those ra text say if I have a tax chunk talking about um like concept A and B and then if I was asking about what's the relation between content concept A and B and then you got lucky you you can it will be very trial to retrieve those content um however if I have a complex set of relations between a large set of uh um uh like entities or or um or like things I care about in that case like we just don't have enough data corpers or enough text um from those documentations which talks about every single aspect of the things so that so that a systematic way of tackling this problem is actually using the knowledge graph to extract those information to establish the connection between those individual Concepts and then together we build a graph uh and you know this graph can be as comp as complex as the the right or it could be as simple as the left well depends on how um the amount of information that your uh application wants to you know leverage um say if if you have you know simple just graph architecture then in that case you can um you can just use this as is which is a already a good way of um encoding those information however if you have a very complex set of um entities and relations relations between them um you can further apply data mining techniques by you know having um you can do clustering of the knowledge entities um and you can extract summary out of it you can even you know um as the graph rag project published by Microsoft you can even establish the layer the structure of the summaries or they they call the communities like um with this kind of technique there is just unlimited ways of um extracting the kind of extracting the the essence of the information from your documents from your uh document corpers uh with this idea um one of the uh one of my favorite work um on this in this domain is called hippo rag um so this research actually tries to um you know combine the um Vector embedding based retrieval which is pretty much just semantic search of the raw document chunks whatever you have and then um combine it with the um the Mind information uh stored in the knowledge graph and then in this work it it does uh some fancy stuff on the uh graph retrieval which is um for each query um it actually it first finds what's the most uh related entity from this graph to my query and then from this starting point it actually runs a page rank so-called personal personalized page rank algorithm to kind of um to flow the in the the the information through the relations and um entities and this structured um uh structured uh uh uh representation and then you know with this just like how page rank in traditional web search helps you find the useful information it uses this algorithm with the same intuition to find the information could potentially help your query however um this is really expensive because it happens every time for for the uh you know for each query you need to run another page Rank and Page rank is no cheap algorithm um so inspired by this we actually did some experimentation and figured um a relatively straightforward but very effective approach uh of combining the you know the power of Knowledge Graph and uh uh Vector search in this rag use case so um how it how it works um so just like like as any search system you need offline indexing phase and online retrieval phase so for the offline processing it actually uh analyze the cors just like how you do a um a so-called convention R conventional rack I mean it has just been here for I think less than 24 months but has been conventional um so as as the traditional rack does it does the chunking of the documents from your data corpers um however um on top of this it also uses large language model to extract the entities and relationships out of it well you can store it in a graph database like uh NE 4J and then do further mining of this and extract more information um so with this analyzed and structured information um you can uh put uh you can you can do uh so-called knowledge embedding of the knowledge entities and relations uh relationships and then store those embedding those Vector rep representation in Vector database like M us and then you can continue this retrieval approach with um you know a slightly um enhanced approach which will I will cover a bit in detail in a bit so um the first problem we need to tackle is after we have the knowledge entities and relationships um extracted from those uh Knowledge Graph how do we actually encode them into Vector embeddings well just like the way of encoding ra text or image and and video and other um any unstructured data we can actually apply a very simple approach which is to uh represent the relation the sorry the triple of two entities and one relation into just a plain uh piece of text and then this this and then with with this text you can use any text embedding model to encode it into Vector this um well just contrary to to our uh to our expectation this actually works pretty good in practice um once you have this you can um go through you know a slightly more complex uh flow of retrieval uh at online serving time so given the query we find the we find the most adjacent entities and relations uh relationships from the vector database and then with this um you know we want to really want to Leverage The Power of the graph so that instead of just using this we'll expand the uh expand the node into is adjacent subgraphs and with this we can rank it by uh large language model to actually filter those things that are that are not really relevant to your search query because by you know expanding it into subgraph you may encounter things that are totally unrelated to your query uh despite the the kind of origin entity was related to your quiry and then with this reun relationships uh we will get the metadata which are the passages where those entities and relations were extracted and then this will be the final um data used for large large language model generation you can put them into your uh prompt and maybe use this spy to optimize that prompt so to see this in action um we actually appli two uh we actually constructed two tables one the relation table which um uh which contains um the text representation of the uh knowledge triple like Source entity relation and destination entity and then we do have the entity table as well which has the um you know the the text representation of the entity which is pretty much the name of the entity and in addition to this we of course we have Vector embedding of either the entity or the relation and uh we also have other metadata Fields so that it can help you to for example U get the um the rout text out where this entity or relationship was extracted from and the related entity so that uh if you are just doing say one half of retrieval you can just directly use this data um so here is how the works um in terms of uh expansion finding the subgraph and a large language model based reranking so first of all we encode the query into a vector eding and then search this vctor embeding in both of the tables with their related um you know entities and relationships and um here the dot uh denotes The Entity and the um the a is actually a relationship um with this we want to expand it a a bit further so that we can know the you know this will help us solve the multihub problem um so we do this by a Brute Force approach we just expand to everything adjacent to it um limited by the number of hops say if we do this by one degree and then we just expand once we want to do it by two degree probably you will get one all one order of magnitude more data but that's okay um we'll see um in the next stab we actually merge those um uh we'll actually merge those uh Ral results and then rank them by large model so let's say here we have um retrieved maybe 30 um relationships and entities um that might be too many for us to you know GA original raw text chunks and then stuff them into the large models prompt because that will just burn your tokens um so instead of um doing that we'll actually do a large language model based reranking by carefully crafting a prompt which pretty much is just says here's the question I want to solve um in this case u i was fascinated by The Crusade um kind of History so what's who is the matter of the leader of the Third Crusade um and and uh when was uh he his or her mother uh born and in this case actually has quite a few hops of information um needed in order to answer this question and was listed be behind this question are all of the um here in this case relationships that we have retrieved before because if there was anything from the document corpers that can help it it's very likely that the those information are spread across the graph and the information are not spread across very distant from each other they will be related I mean this is just a plain observation of how human like how how how we organize or knowledges um so with this intuition um we we know that well from this list there will be probably enough information to answer this or otherwise you just simply don't have this information from your document corers then nothing can help other than you uh you kind of construct those documentations which is also a very useful technique to um to improve the ritual quality even though it sounds like a bit too plain um that being said with those relationss you can see like some of the uh information are not really related to the or not helpful to answering this question uh however there are some in this case the lar grid model um ex uses the um it B basically just think out loud so it's very helpful to force the large model to speak out loud regarding his thought process so in this case we we use this prompt technique and it actually explains okay to to find the answer of this question I need to reason about the relation between those con Concepts and the result of that reasoning is that it figured three uh pieces of information in this case those are relationships from this knowledge graph where the rich Richard the the lonard actually led the Third Crusade and who is the mother of Richard and like this this woman was born in which year so with this we are really close to solving this whole problem but just one last piece um there's actually another research that analyzed um the effect of you know the results of you solely using the text representation of the knowledge entities and relations versus directly using the raw document chunks from your your data corers I mean apparently if you're using the document chunks you're consuming more tokens in the large the inference phase but this is actually useful because remember we actually used large langage model or a mechanical way to extract the knowledge entities this is not always reliable I mean it works pretty well in practice but it's not always reliable especially compared to having human beings uh you know annotating this by hand annotating the um the like or constructing the knowledge graph by you know just looking at the text you can H you can hide a a lot of um uh labor to do this but that will be pretty expensive so um with this compromise it's actually very meaningful to uh directly use the ra passage where you extracted those entities well you can also you can you can you can do more by you know stuffing the text representation of the entities there as well U we didn't test this but I think it could like it could provide a a little bit of help um but maybe it's not worth the the xer ha because you with those ra text chunks you have already have enough you know materials to cover those information represented by the knowledge graph uh by the entities from the knowledge graph um so let's look at the the final result We compare this with uh both naive Rag and which is the blue bars uh sorry which is the green bar and hippo which used the personalized page rank algorithm where for each query you need to anchor on a node and then run the um is kind of um a a walkthrough or traversal algorithm uh through the whole subgraph or or even the the entire graph so compared to that the um the straightforward approach of using um only uh knowledge embedding vectors for uh retrieval in the online serving phase actually does better than uh both hippo Rag and N rag um so if you would like to check more uh if you would like to um you know uh have more fun and and play with this approach by yourself um we actually have the open source m u that supports all of the deployment most um from as simple as you know installing Ms Light in your laptop or notebook by just PP install to very performant and scalable Docker and kubernetes deployment and also uh if you are using the um if you are um um if you favor uh using uh API service we also have z z is cloud which is U managed M um uh by us and we provide Ser both seres and dedicated instance for you to get started with and the M project because both of the um self- deployed M and um the commercial zus Cloud offering um they share the same API so that you you pretty much just need to develop the application logic once and then you can use it anywhere with whatever environment and they all work with a whole lot of um Frameworks in the rag and uh AI application uh development space including dispy and other projects so um that's all from my uh sharing um so if you like to check more about M here's the QR code for the uh for good GitHub page and we also have the Z Cloud uh the fully managed MERS um as a s offering so um if you are interested in this kind of work um like tutorials and best practices in applying Vector search in your gen applications uh we have more tutorials and demonstrations uh from me .io uh we have you know multimodel rag uh uh graph Rag and other like hybrid search all those um uh examples that uh will be you know helpful in uh if you are exploring the semantic search space so that's it um ready for classs so we're a bit out of time but I think questions okay question or yeah um there are actually more uh complex ways of creating embeding for knowledge entities for example um um there are approaches where they encode the subjects the the entities and the relations and then with the even the rout uh do chunks where those entities were extracted from um we tried that but um one interesting observation is that a um just simply using the the text representation just concatenating the the names of the entities and relations that works pretty well in in practice um you know in terms of beating the um uh both the traditional Rag and uh uh personalized page rank based uh graph graph plus uh sorry kg PL R approach that's one thing so um um this approach is more reliable and more widely available because there are just tens of tax in beding models but there are not as many um graph uh sorry knowledge entity in beding models um on the other hand is we observed uh Improvement on the overall search quality by simply upgrading the eding model the tax eding model from a smaller one to a larger one so that if that's really the you know the um the approach uh the direction of optimization you want to go after um just simply using more advanced more recent eding models can provide um you know SI sizable improvements in terms of the uh uh um summarized score of the ritual quality um as as at least uh evaluated by those public data sites so I think um definitely there there is um we we do have interest and I think there there is a need to explore more advanced graph eding techniques uh um but in this case um there are also other places to to optimize such as the um performance and latency of the retrieval uh approach because it it does involve large remodel as part of the retrieval which is used for ranking um and also we tried um you know expanding it Beyond two HS like we tried expanding to three and four Hops and then doing a triming of the things um which from that without using large langr model and we found um interesting results that you know simply doing two half of retrieval that gives you a um a good enough results and expanding towards say three to four hops it doesn't really help that much so I I think um like those kind of optimizations are maybe they they they are more like uh cost effective than exploring the knowledge embedding because you know the current um landscape of embeding models are like most of the embeding models they are based on bird or large langry model architecture and some you know even if you can invent a eding architecture eding model architecture that kind of optimize for the knowledge entities if there's not enough attention from the community and from the industry um you don't have enough you know good training data to imp improve its performance it's very likely that it cannot still compete with the this very simple approach thank you app [Applause]