Graph Exchange: Rishi Puri, GNNs & LLMs in PyG

Graph Exchange: Rishi Puri, GNNs & LLMs in PyG

Recording: Graph Exchange: Rishi Puri, GNNs & LLMs in PyG

So hello everybody. I'm a senior deep learning engineer at NVIDIA and lead our engineering efforts for PyTorch Geometric and research efforts for combining graph and language modeling. Today I will be discussing GNN's and LLMs in PIG. Modern LLMs are built on the transformer model and excel at predicting the next token in a sequence. They are based on the attention module introduced in the famous attention is all you need paper from 2017. Since then, we have seen tremendous progress in LLM research simply by scaling up model size in step with data set size. However, we are reaching the limits of these famous scaling laws and are seeing a plateau in accuracy across the industry when it comes to simple auto reggressive pre-training. Now, we are seeing ways to break the plateau with test time chain of thought scaling

But due to the limited time in this session, I will not go into this area of research. But rest assured, we are working on enhancing this with GNN's as well. So stay tuned to my LinkedIn for future announcements about this. So, one issue with auto reggressively pre-trained LLMs is that their underlying predictive task is to predict the next token given previous contextual tokens. The problem is that this method of computational thinking is inherently single hop in nature which is far more simplistic than what humans do naturally. One issue with autog uh sorry this single hop logic it then forms a sort of associative memory where given context they produce the most probabilistically relevant output tokens based on the distribution learned across their training data. As we are seeing today basic LLMs as a single hop memory bank are not sufficient for answering out of distribution questions reliably. This is essential problem to solve since there is no limit to edge cases for trans areas where transformers are used like NLP, robotics, medicine and more

Imagine your goal is to learn the sequence 2 4 6 8 and so on. Our general human intelligence can easily identify the solution that will handle the entire distribution of data just plus two. But an associate memory model such as basic auto reggressive LLMs will likely fail on out of distribution examples. Imagine for training we have 500 GB of space to store our model's parameters. We could just directly store the next token in the sequence for every available token. This is what LLMs do today with gradient descentbased learning. The parameters will basically just store 2 to 4, 4 to 6 and so on. When we are testing this will pass in the trained distribution but will generally fail when we are out of distribution

the correct output is not well defined for the model and in the case of LLMs this causes them to hallucinate. The point of these examples was to showcase how simple LMS are much closer to an associative memory than general human intelligence. This is a slide from a recent talk given by one of the founding fathers of modern AI, Yan Lun. See the QR code in the corner for the slow full slide deck that this comes from. In it, he brings up one of the key issues underlying basic LLM today. LMS learn the distribution of token n given token 0 to n minus one. Then when you use them for inference, they generate the most likely next token given the prompt. Then they combine this with the prompt and repeat this process over and over again

This is done in a loop until a stop token is generated or a user selected maximum length is reached. The key point of lacun's slides that I want to highlight is a part mentioning the probability of being correct is a diverging function with respect to the number of tokens n for the equation in lun slide. E represents the error for next token prediction. Even if the error for next token prediction is near zero, the chance of our entire sequence being wrong becomes high very quickly. We will circle back to dive deeper into this issue shortly. This slide has been adopted from a 2023 talk given by the creator of PIG, Matias Feay. He defines GNN's as models that update node representations by repeatedly transforming and aggregating representations of direct neighbors. The diagrams here explain this quite well and I'll walk through each of them

The top is the most most basic example. We have a single layer of a GNN which computes the embeddings of a single node in this case node one as a function of the embeddings of its direct onehop neighbors. One can think of convolutional neural networks and transformers of sub as subclasses of GNN's much the way that triangles and quadrilaterals can be thought of as subclasses of polygons. The diagram for the convenet shows how a convenet across pixels is the same as a GNN over a graph of pixel nodes connected to their direct neighbors in the image. For the transformer diagram, you can see how the transformer attention layer is essentially a GNN over a fully connected graph of the tokens. In general, an end layer GNN learns to predict over n hops of your graph, which makes them excellent candidates for enabling multihop reasoning in AI systems. This is essential since multihop reasoning is a key aspect of human intelligence. An example of my two-h neighborhood is me, my friends, and all of their friends

Our research finds that GNN's can improve LM accuracy. We believe this to be possible since GNN embeddings on graph data can provide a semi-orththogonal source of information. LMS also tend to be much larger than GNN's. Remember, a transformer layer is essentially a GNN on a fully connected graph of tokens. However, the transformer uses parameters to learn the importance of each connection. This is what attention is. While GNN's do not need to parameterize how the entities are connected since this is provided by the graph data. As a consequence, adding in GNN's has almost no effect on the total number of parameters

For example, one could greatly improve the accuracy of an 8 billion parameter LLM with a GNN as small as 10 million parameters. For any given input, the output is deterministic for most GNN's. For LLMs, this is not the case. They are generative AI. This means that they learn a distribution and then if you want to use them for inference, you randomly generate outputs that align with that distribution. Imagine we have a prompt with two or more valid answers. If both are present in the training data, this can lead to a high chance of plausible sounding hallucinations. As a toy example, imagine we train an LLM to predict the next token given two training documents

One with the sequence 2468 and the other with a sequence 261854. Now, what would the LLM predict as the next token for two? It's unclear. In this case, four and six are both equally likely outputs. Now, to relate this to the real world, imagine we train on documents that contain task descriptions and then their corresponding Python solution. If we were to train on data from leak code, a popular coding prep website, we could would have many documents with the same task but different solutions that are all valid at inference time. This will often result in a in the model producing a solution that looks valid at first glance but has sneaky issues. One can imagine why by thinking about what would happen if they looked at two correct code solutions online and then copy and pasted random chunks from each of them into their own code without understanding what they do. This will rarely result in high functioning code and you've probably seen people do this in college and it doesn't work out well for them

Um, this is a huge problem for coding projects and it often results in devs having to spend substantial amounts of time debugging their AI generated code and it can also allow it can also allow for extremely damaging bugs to slip through. Some examples are fatal errors that result in crashes or allowing bad actors to hack into your system. Basic LLMs are great for creative tasks that don't need absolute precision. Some examples are fields like writing and art. Your feed has probably been flooded with Yeah. So, impressive art from chat GBT, you know, the studio giveway trend, and there's countless others that are very awesome. But LMS are not ideal for tasks that require extreme precision like coding or making high impact decisions. In these cases, even minor hallucinations can result in extremely damaging outcomes

For example, I recently saw an article where lawyers are being sued for using AI that creates false citations. I saw one even more recently this morning that hikers are needing search and rescue because they just blindly follow chat GBT into the wild. Um, you can imagine chain of thought, reasoning, and test time scaling can help in cases where you have verifiable rewards like math and coding, but there's not really a verifiable reward for cases like this, right? So, we still have a problem to solve. Additionally, we are running out of public data to train our LLMs in the hopes of solving these problems. Graph models allow us to add a whole new realm of data for training frontier models. While we do not believe that GNN's will solve all of the problems I have presented today, they are a step in the right direction. So, here we have a diagram of a madeup Nvidia org chart. The question is, who is the shared co-orker between Bob Joe and Ronald Kohiman? At first glance, the answer may not be perfectly obvious even to a human

But thanks to the highlighting of the key nodes and edges, the answer becomes obvious. Arno source finger. Using PIG, a GNN plus LLM gets the answer right, while the corresponding LLM fine-tuned the same way fails to answer correctly, even though both models were fine-tuned on the same general knowledge toy data set. The pure LLM approach I described is generally similar to most graph rag workflows today, such as the one popularized from Microsoft. At a high level, I like to describe the GNN's role as highlighting the relevant nodes and edges for the LLM, improving upon these basic pure LLM graph rag pipelines. So here we have a diagram that depicts a general graph-based rag workflow using GNN's as a prefix to the LLM's inputs. We start with a natural language query. This needs to be encoded by our LLM's tokenizer and encoder

In general, we know LLM suffer from N squared complexity where N is the number of tokens. We also need to use the query to retrieve a relevant contextual subgraph from a larger knowledge graph database and then encode it with a GNN. The retrieval step is a key part of any rag pipeline. I will go into further detail after this diagram including the computational complexity. The computation for the GNN can be considered inconsequential given that it can have a thousand times less parameters than the LLM and the GNN only has to perform inference on a subset of a few dozen nodes. In detail, we also pass a textified version of the retrieved subgraph to the LLM encoder. Note that most graph rag solutions do not use GNN's and only provide a textified subgraph to the LLM. As you saw before in my toy example, LLMs and even humans can struggle to understand these textified graph formats

This is why our work of adding in GNN's to the pipeline is so crucial as it enables the LLM to understand graphs. Note that we consider passing the textified graph to be an important step since the GNN embedding reduces information down into a single subgraph token. One can view the passing of the textified graph as a skip connection similar to those in the famous ResNet model. In general, skip connections are important to preserve information as it is reduced by the layers in a neural network. We then want to combine the GNN and LLM embeddings and this step is basically computationally inconsequential. Lastly, we pass this into the LLM's decoder module to get a final response. In our work, we generally fine-tune the GNN and the LLM together. However, if you're constrained by compute resources, you can use Laura for the L

You can use Laura for the LLM or you can even freeze the LM. In the case of a frozen LLM, the GNN essentially learns to embed the context subgraph into the latent space of the frozen language model. In general, Laura will be better than frozen and fully fine-tuning the system is the best. Pi expects a graph where each node and possibly each edge has a corresponding feature vector. When working with a textual graph, we normally store this as a triplet of short strings. An example is the triple cats eat dogs. To convert this into a form usable to pi, we need a model that converts strings into feature vectors. Pi sentence transformer takes in a list of strings and outputs a tensor where each vector represents the corresponding input string

This is essential for converting a textual graph into a pi object that we can do retrieval and machine learning with. We need to call the underlying model three times for each edge in our graph. This is because for each edge we must call the model on both entities as well as the relation. Since the model must be called so often we we use a small language model like modern bird for efficiency. Small language models are sufficient for this task since they generally tend to understand short phrases. As an anecdote, when I was in school, we would use now ancient models like glove and LSTM and could inspect that their learned similarities were relatively accurate for short phrases. We generally only need large language models for large or complex bodies of text. As future work, we intend to measure the trade-offs for small versus large language models

Um, and most recently, we've added in the vision transformer module, which allows you to embed images as well for node and edge features. So, this slide zooms into the retrieval part of the previous diagram. If we start at the left side of this slide, you have a knowledge graph source of some kind which can have millions or billions of nodes. To sample a subgraph from it, we use k nearest neighbors between the query and the node embeddings followed by n hop neighbor sampling. After this, you're left with likely thousands of nodes. Next, to trim down to the dozens of most relevant nodes, we apply the prize collecting Steiner tree algorithm or PCST. This method is essentially a semantic similaritybased graph traversal technique. This step is performed on thousands of nodes and has linear scaling

So we can consider it to be computationally free while it brings major improvements in rag accuracy. Our retrieval framework is implemented in PIG in a customizable way to support extensions to domain specific retrieval methods. For example, we worked with Neo4j and then they made their own custom version. So the latest release of PIG comes with new GNN plus LLM features as one of the main updates. There are three new major GNN plus LLM features in PIG. First is the G retriever model that allows a user to easily combine any PIGGNN with any hugging face LM for graph-based rag tasks. The next feature is the web QSP data set for rag question answering with knowledge graph context. The next the final feature is an example of how to train an a G retriever model on the web QSP data set

By default, this uses an Llama 3 8B LLM with a graph attention transformer GNN. There's also a tiny Llama flag for using a 1 billion parameter model, but you can easily swap in any hugging face LLM. In general, GNN plus LLM workloads are very computensive and therefore GPUs are essentially a necessity. running on CPU will basically never finish. So in this CA case study, Neo4j used Stark Prime, a medical knowledge graph data set made by UR's Stanford group. Uh Neo4j was able to use G- Retriever to double the hit at one metric. The previous state-of-the-art is a pure LLM approach. So this shows that adding GNN's to LLMs has immense value when used on a high quality knowledge graph

The baseline of comparison is a 16% accurate 7 billion parameter LLMbased agentic graph rag while the new GNN plus LLM solution provides 32% accuracy. The new solution consists of a llama 318B model using Laura with a 10 million parameter GAT or graph attention transformer. As I mentioned before, this solution shows how much accuracy can be improved by adding a GNN for such inconsequential parametric cost. They have open sourced their work on GitHub and we also wrote a blog about it together. Most rag data sets only have unstructured text context. As such, there is often an additional task where we must convert unstructured text into a knowledge graph. Our existing PIG framework uses the simple KG format of triples of entity, relation, and entity. As mentioned before, LMS are specialized for handling unstructured text

So they are ideal models for this task. As our solution to this problem, we have created the text to KG class in PIG. See the PR link for details. For knowledge graph-based rag, the knowledge graph is our source of information and as such the quality of it is essential. Because of this, we chose the latest open-source model from NVIDIA serviced through NVIDIA inference microservices or NIMS. I chose to use NIMS since most PIG users can't run a 70 billion parameter LLM, but this model was chosen since it's on par with GPT40, but it's open source and much smaller. See the link for details about this model. As a rough measurement, you can assume this process incurs an upfront cost of about 2 milliseconds per character from the docs that are being parsed into a kg

This is with the default NIM subscription, but you can obviously scale up or down based on your usage needs. Um, there is also an option for using a local language model. For this, the smallest language model I could get to work was 14 billion parameters. The model was chosen because it was the best 14 billion parameter model on HuggingFace leaderboard at the time. If you look at the bottom left, that snippet is the current system prompt for this class. The snippet on the right shows how easy it is to use. In general, there are cases where vector rag alone does well and there are cases where graph rag struggles to shine in comparison, especially when the correct information needed to answer is found in a single document. Graph rag generally shines when the correct information is spread across multiple documents

However, we should expect that combining graph and vector rag is always at least as good or likely better than either alone. GNN's help the LLM decoder to understand graphs, producing much better accuracy. This diagram covers the fully combined workflow. As you can see, this is the same diagram from before, but I have added in a block to insert the documents retrieved by any sort of document retrieval method. Since our system is fine-tuned, the LLM should learn to reason over both sources of information. Even if you freeze the LLM, it should, you know, be able to handle the retrieve documents and the GNN should learn how to speak its language. In our implementation, we simply inject the retrieve text documents at the end of our prompt and we have the textified version of our retrieved graph prepending the prompt. We find this to be the optimal solution for now and we'll continue to refine this as we explore further

So this slide covers a demo made by our PM Sentosh Pavani. He was able to upload our Pygi code to cursor and have cursor completely vive code a graphical interface showing how easy our framework is to build upon. I'll skip the demo for now, but I'll come back to it if time allows. The idea of using GNN embeddings as a prefix to transformer sequence models is highly general and can be extended to many modalities. The goal of the latest PIG community sprint was to expand GNN plus LM features to support the sciences like biology and chemistry. Check out the linked GitHub issue which tracks the community effort. There are four projects going on to integrate techniques from three biology papers and one chemistry paper using GNN's and LLMs. The general goal of these projects is to advance medicine and science and integrate general reusable APIs for GNN plus LM workflows

As an example, the vision transformer module I mentioned before was added as part of this sprint. Three out of four of the tasks are completed so far. Check them out in the examples on the PIG GitHub or the NVIDIA container. The possibilities are truly endless. As another example, one could imagine a setting where a business leader could ask an AI general questions about their customers. You could use this for any set of docs that have links referring to each other, like those in Wikipedia, academic journals, or coding docs to name a few examples. This idea can be further extended to use multiple modalities for a single task. Imagine an Amazon product graph where each node has a text review and a photo or hyperlink documents with both text and images

Extending even further, one could use highly heterogeneous graphs such as those in Realbench or Kumo where node and edge features could be text, images, audio, and more. Basically, anything that can be passed through a GNN in Py could be integrated into this framework. We actually have an interesting spectrum of graph-like data. On the left side of the spectrum, we have the setting of unstructured data in the form of documents that we must parse into a knowledge graph. This setting is inherently less accurate since we rely on an AI model for this parsing. On the other side of the spectrum, we have data like molecules, hyperlink document graphs, or relational table data like in real bench and kumo. In this setting, we can assume our graph data is highly accurate. In general, we know that most of these sources will cover different knowledge and we also know that most enterprises have multiple of these data sources

Using PIG, you could set up a separate GNN for each type of graph data source, all providing context to a master LLM following similar patterns as discussed before. As a final note, today there are two kinds of GNN's. The traditional kind are the message passing GNN's such as GAT which I mentioned before. And then there are graph transformers which are newer. See the QR code for a webinar on the topic I gave with Stanford and Kumo. I'm working with the community to get more graph transformers added to PIG. My goal would be to see how these affect the accuracy of the previously discussed GNN plus LLM systems. In conclusion, this work is built and optimized for the NVIDIA PIG container

Also, look out for f future productization coming soon. Also, in case any details were missed, check out the recorded webinar I did with Kumo a few weeks back. It's on YouTube. It's at the QR code. As a final note, I want to acknowledge my fellow Invidians who have contributed to this effort and all of the PiG contributors. The PIG community is ever growing and contributions are more than welcome. Thank you. [Applause]