LLMs + GNNs

This is a guest lecture by NVIDIA's PyG (PyTorch Geometric) team. It starts with a result, to show why combining GNNs with LLMs is worth the effort.

Slide 4. Neo4j is a graph database company (it stores data as nodes and relationships and queries them with the Cypher language). They used PyG 2.6's G-Retriever (explained later in this lecture) on STaRK-Prime, a question-answering benchmark built on a biomedical "semi-structured knowledge base". The figure at the bottom shows that knowledge base: node types such as Gene, Drug, Disease, Pathway, Protein, Anatomy, Phenotype and Molecular function, linked by relations like "carrier", "interact with", "indication/contraindication", "expression present/absent" and "associated with". Each node also carries free text, for example a description of the gene GPANK1 or of the disease GM1 gangliosidosis type I. So the answer to a question can depend both on text and on graph structure.

• Result: Hits@1 doubled, from 0.16 to 0.32. The 0.16 baseline is an agentic GraphRAG system using a 7B LLM (an LLM that issues graph queries step by step); the 0.32 system is the PyG GNN+LLM solution.
• The PyG solution is small: Llama 3.1-8B fine-tuned with LoRA plus a 10M-parameter GAT.
• It even beats agentic GraphRAG built on much larger frontier models: Claude-3-Opus (0.18) and GPT-4-Turbo (0.20).

Recall that Hits@1 is the fraction of questions where the top-ranked answer is correct. The message of the slide: adding a small trained GNN that reads the graph helped more than using a bigger LLM that only "talks to" the graph through text queries.

Slide 5 reviews the transformer (the figure is the encoder–decoder diagram from "Attention Is All You Need": input and output embeddings with positional encoding, stacked multi-head attention and feed-forward blocks, then a linear layer and softmax giving output probabilities). The speakers' claims:

• LLMs are excellent at predicting the next token in a sequence.
• In their view, progress from "attention is all you need" alone is plateauing.
• LLMs are pretrained for single-hop logic: given context tokens, predict the next tokens. The speakers call this associative memory: it works in distribution but fails out of distribution, and NLP, robotics, medicine and other fields contain infinitely many out-of-distribution tasks.

Slide 6 gives a toy picture. Suppose the goal is to learn the sequence "2, 4, 6, 8, 10, 12, 14, …".

• General (human) intelligence learns the rule "given a number, return that number + 2". This never fails out of distribution.
• An associative memory model instead stores, for 500 GB of numbers, the next token: 2→4, 4→6, 6→8, … At test time it passes on numbers it has seen and fails on new ones: the output is undefined, so it hallucinates.

The speakers' claim is that today's LLMs are closer to associative memory than to human intelligence. (This is a position, not a theorem; treat it as the motivation of the talk.)

Slide 7 quotes Yann LeCun's slide "Auto-regressive generative models suck!". The argument: let \(e\) be the probability that any produced token takes the answer outside the set of correct answers. If errors are independent, the probability that an answer of length \(n\) is correct is

which decays exponentially in \(n\). The pie chart shows the "tree of all possible token sequences" as a large disk and the "subtree of correct answers" as a thin red slice: one wrong token can push you out of the slice for good. LeCun concludes that auto-regressive LLMs cannot be made factual or controllable without a major redesign (see also Dziri et al., arXiv:2305.18654).

The independence assumption is strong: real models can correct themselves, and errors are correlated. The formula is an intuition for why long, precise outputs are fragile, not an exact law.

Slide 8 recalls what a GNN is (figure from Matthias Fey's 2023 Graph Learning Conference talk, plus a QR code to the "GNN 101" visualizer from the University of Minnesota). GNNs update node representations by repeatedly transforming and aggregating representations of direct neighbors. The neural message passing scheme computes, for a node, messages \(\text{Message}(h_i^{(\ell)}, h_j^{(\ell)}, e_{ij}^{(\ell)})\) from each neighbor \(j\) and aggregates them. This gives data-dependent computation graphs, and generalizes architectures from CNNs (grids) to Transformers (complete graphs) to general GNNs.

Slide 9. Three points:

• An \(N\)-layer GNN performs \(N\) hops of "logic": each layer passes information one hop. The picture shows node A's 2-hop computation graph (A ← B, C, D, E; those ← their own neighbors).
• GNNs can improve LLM accuracy because they carry semi-orthogonal information: structure that the text alone does not show.
• \(\text{size(LLM)} \gg \text{size(GNN)}\). Quoting Jure Leskovec, "a Transformer layer can be seen as a special GNN that runs on a fully connected 'word' graph". So adding a GNN to an LLM system adds almost no parameters.

Slide 10: GNNs are deterministic, LLMs are not. For a given input, most GNNs output the same answer every time. LLMs are generative: they learn a distribution and then sample outputs that fit it. One prompt can have two or more "right" answers, so there is a high chance of a plausible-sounding hallucination. Toy example: train an LLM on two sequences, 2, 4, 6, 8 and 2, 6, 18, 54. What should it predict after "2"? Both 4 (add 2) and 6 (multiply by 3) fit the training data.

Slide 11: the real-world version. For a task plus a Python solution there are almost always many valid solutions. An LLM can produce output that looks valid but has issues, because combining random chunks from correct codebases often does not give a correct program. Consequences: developers using AI spend a lot of time debugging, and severe bugs can slip through (crashes, security holes).

Slide 12: strengths and weaknesses. Basic LLMs are great for creative tasks that do not need absolute precision (writing, art). They are not ideal for precise tasks such as coding or high-impact decisions. Also, LLMs are running out of public data; GNNs open a whole new world of data (relational, structured, often private). GNNs will not solve all problems, but they are a step in the right direction.

Slide 13: VectorRAG vs GraphRAG. RAG = Retrieval-Augmented Generation: before answering, retrieve relevant context and give it to the LLM.

• VectorRAG: embed all documents; retrieve the top-\(K\) documents whose embedding vectors are most similar to the query. Good enough when the answer is inside a single document.
• GraphRAG: retrieve a relevant subgraph. Good when the answer needs several documents with related entities, that is, multi-hop reasoning. The diagram previews the G-Retriever pipeline used later.

Slide 14: a made-up NVIDIA org chart. Ten people (nodes 0–9: Ron Wonald, Jensen Huang, Ronald Coleyman, Arno Swarzfinger, Bob Joe, Joe Bob, Richard Preor, Breance Owls, Justin Brogard, Ash Katchoop) and typed edges such as "reports.to", "leads", "mentors" and "works.with". The graph is given to the model as two text tables, node_id,node_attr and src,edge_attr,dst. Question: "Who is the shared coworker between Bob Joe and Ronald Coleyman?" The answer is Arno Swarzfinger, because of the edges 4,works.with,3 (Bob Joe – Arno) and 3,works.with,2 (Arno – Ronald Coleyman). The slide reports: LLM alone is wrong; GNN+LLM is correct.

Why is this hard for an LLM? The two needed edges are scattered among the edge list, and the model must join them on the shared node id 3 and map ids back to names: a 2-hop query. The figure also has distractors with similar names (Bob Joe vs Joe Bob, Ron Wonald vs Ronald Coleyman). A 2-layer GNN does exactly this kind of 2-hop join.

Slide 15: a history lesson. "Attention Is All You Need" appeared on arXiv on June 12, 2017; the GraphSAGE paper appeared on June 7, 2017, five days earlier. Since then industry has put most of its effort into LLMs, while a small but strong PyG community grew in academia, plus a small group of companies using GNNs (Pinterest, Kumo, Spotify, …). PyG's GNN+LLM work aims to combine the two lines and move AI systems past the limits of LLMs.

Slides 16–21 build up one diagram, adding one step per slide (the red box marks the new part). This is G-Retriever (arXiv:2402.07630), a GraphRAG method where the GNN feeds the LLM. Given a natural-language query at inference time:

1. Encode the query (slide 17): tokenize the query and encode it with the LLM encoder. Cost: LLM complexity is \(O(n^2)\) in the number of tokens \(n\), because attention compares all token pairs.
2. Retrieve and encode a subgraph (slide 18): retrieve a subgraph of the knowledge graph (KG) relevant to the query, take its node/edge features, and encode it with a GNN encoder. The slide says GNN complexity is "~free" compared with the LLM. (Node/edge embeddings are covered on slides 23–25 and the retrieval algorithm on slide 26.)
3. (2a, slide 19) Optionally also turn the retrieved subgraph into text (the textified graph, like the tables on slide 14) and pass it into the LLM encoder with the query.
4. Jointly embed (slide 20): concatenate the GNN embedding with the LLM token embeddings. In practice the graph embedding is projected to the LLM's embedding size and placed in front of the prompt, like a "soft prompt" token.
5. Decode (slide 21): the LLM decoder reads the joint embedding and generates the natural-language response.

Slide 22: training options, best to worst test accuracy: (1) full fine-tuning of the LLM, (2) LoRA fine-tuning of the LLM (next best), (3) frozen LLM (least good). In all cases the GNN (and projection) are trained, so gradients from the answer tokens flow back into the GNN.

Intuition: the textified graph gives the LLM the raw facts, but the LLM must still do multi-hop joins in its \(O(n^2)\) attention. The GNN embedding gives it a pre-digested, structure-aware summary that costs almost nothing.

Slide 23. PyG expects a feature vector for each node and, optionally, for each edge. But in a text KG, entities and relations are short phrases. Example triple: (cats, eat, dogs). So we need a model \(\text{Model}(\text{str}) \to \text{feature vector}\).

Slide 24: sentence transformer. A sentence transformer maps a list of strings to an embedding tensor: \(\text{model}(\text{List[str]}) \to \text{Tensor}\). The code on the slide loads SentenceTransformer(model_name='sentence-transformers/all-roberta-large-v1') and moves it to the GPU if available. Cost: the model is called \(3 \times \text{num\_edges}\) times, because each edge needs the head entity, the relation and the tail entity embedded. So use a small language model (SLM) such as ModernBERT for efficiency.

Slide 25: why small LMs are fine here. Small LMs understand short phrases well enough; large LMs are only needed for long or complex text. Future work: measure the trade-off between SLM and LLM encoders. The latest PyG also supports Vision Transformers, so image attributes can be encoded the same way.

In practice you embed each distinct entity and relation string once and reuse it, so the real cost is closer to the number of distinct strings than to \(3\times\) the number of edges.

Slide 26 zooms into the "Retrieve Subgraph Context" box of the G-Retriever diagram. The input is a remote KG with millions of nodes and the prompt. Retrieval runs in two stages:

1. k-NN and n-hop neighbors. Embed the query; find the \(K\) nodes (and edges) most similar to it; expand each seed by \(n\) hops. This gives an untrimmed subgraph of thousands of nodes. Cost: with \(X\) seed nodes, \(n\) hops and fanout \(f\), it is \(O(X \cdot K \cdot f^n)\), exponential in the number of hops, like neighbor sampling in GraphSAGE.
2. Prize-Collecting Steiner Tree (PCST). Prune the subgraph to a connected subgraph of dozens of nodes. The slide calls it "similarity-based graph traversal"; it is linear in the thousands of nodes, so "free" compared with the LLM.

What PCST does (background from the G-Retriever paper, not on the slide): each node and edge gets a prize that is larger when it is more similar to the query; each edge also has a cost. PCST finds a connected subgraph \(S\) that maximizes

Connectedness matters: a relevant entity is only useful together with the path that links it to the other relevant entities. Pure top-\(K\) retrieval can return relevant but disconnected fragments; PCST keeps the "bridges" even when a bridge node itself is not very similar to the query.

Slide 27 contrasts the two directions:

• GNN feeds LLM. Data: a prompt conditioned on a graph. Tasks: token-level tasks, for example question answering. Flow: Graph → GNN = GNN_out; then (GNN_out + prompt) → LLM = answer tokens.
• LLM feeds GNN. Data: a graph where every node/edge has tokens (text) attached. Tasks: graph learning, for example node classification or link prediction. Flow: compute a feature vector for each node/edge with a sentence or vision transformer, then feed the graph to a GNN.

Slide 28: GLEM (arXiv:2210.14709) is described as the state of the art for node classification on text-attributed graphs (TAGs), for example ogbn-products (Amazon products whose descriptions are text). The figure: raw text of node objects is tokenized into a text-attributed graph. Then an EM loop runs:

• E-step (yellow): train the language model on the text of each node; it outputs predictions used as pseudo-labels.
• M-step (blue): train the GNN on the graph; for the M-step, the LM's pseudo-outputs are concatenated with the original node features. The GNN's predictions become pseudo-labels for the next E-step.
• Pseudo-labels from each model serve as augmented training data for the other; the final output is the external prediction.

Why alternate instead of training one big LM+GNN end to end? Back-propagating through a large LM for every node in a big graph is too expensive. Alternating lets each model learn from the other: the LM learns what the graph "thinks" about a node, and the GNN gets text-aware inputs. Unlabeled nodes still contribute through pseudo-labels, which is a form of semi-supervised learning (compare label propagation in Chapter 19 of the notes).

Slide 29: GLEM in PyG. The PyG implementation is 7× faster than the original paper's code on ogbn-products. It comes with a new TAG (text-attributed graph) interface, optional LLM fine-tuning with LoRA (via Hugging Face's PEFT library), and an example at examples/llm/glem.py.

Slide 30: TAG usage. Load PygNodePropPredDataset('ogbn-products'), get the split, and wrap it: TAGDataset(root, dataset, "prajjwal1/bert-tiny", token_on_disk=...). It takes a path, the dataset and the LM of your choice; optionally it saves tokens on disk.

Slide 31: GLEM setup. Build a GNN, for example GAT(in_channels=1024, hidden_channels=1024, num_layers=4, out_channels=num_classes, heads=4), then GLEM(lm_to_use='prajjwal1/bert-tiny', gnn_to_use=gnn, out_channels=num_classes, lm_use_lora=..., device=device). Options: use LoRA; choose the device (the default is CPU, which is very slow).

Slide 32: setup. A G-Retriever model is an LLM plus a GNN:

• gnn = GAT(in_channels=1024, hidden_channels=1024, out_channels=1024, num_layers=4, heads=4). The 1024 matches the sentence-transformer feature size; the output must later be projected to the LLM's hidden size.
• llm = LLM(model_name="meta-llama/Meta-Llama-3.1-8B-Instruct")
• model = GRetriever(llm=llm, gnn=gnn). GPU setup is automatic.

Slide 33: getting the loss. Calling model(...) in training mode returns the loss. Its arguments, one mini-batch of questions:

• batch.question: list of question strings;
• batch.x: node features, shape [num_nodes, num_features];
• batch.edge_index: shape [2, num_edges];
• batch.batch: which nodes belong to which question (each question has its own retrieved subgraph, packed into one big disconnected graph, as in PyG graph batching);
• batch.label: list of answers (the target text);
• batch.edge_attr: edge features, [num_edges, num_features];
• batch.desc: list of text descriptions of each graph (the textified graph of step 2a).

The loss is the usual next-token cross-entropy on the answer tokens, given the graph embedding plus the question.

Slide 34: inference. model.inference(questions, batch.x, batch.edge_index, batch.batch, batch.edge_attr, textified_graphs): the same inputs except the labels; edge attributes and textified graphs are optional but recommended.

Slide 35: knowledge graph creation. Most RAG datasets only have unstructured text. To use GraphRAG we need a KG, so the task is: unstructured text → KG in the triple format (entity_1, relation, entity_2). LLMs are specialized for unstructured text, so they are the ideal tool for this (the picture: documents → LLM → graph).

Slide 36: the TXT2KG class, released in PyG 2.7. It is prompt engineering plus Python filters. The system prompt asks the model to convert the text into a list of knowledge triples of the form ('entity', 'relation', 'entity'), separated by new lines, to output nothing else, to focus on key triples, and to reuse the same entities whenever possible (so that the triples link into one graph rather than many duplicates). Usage: loop over documents and call kg_maker.add_doc_2_KG(txt=context_doc).

Slide 37: GraphRAG + VectorRAG = PyG-RAG (torch_geometric.llm). The diagram is the G-Retriever pipeline with two additions: at the top, documents are turned into the KG by an LLM (TXT2KG); at the bottom, a Doc/Chunk Retrieval box (classic VectorRAG) feeds retrieved text chunks into the LLM encoder alongside the query and the textified graph. So the LLM gets three inputs: the GNN embedding, the textified subgraph and the raw text chunks. The slide says PyG-RAG is being adopted across Fortune 500 companies and inside NVIDIA; example code: examples/llm/txt2kg_rag.py.

Slide 38: accuracy on multi-hop synthetic Q&A for a lecture corpus. The bar chart shows accuracy judged by a 253B Nemotron-Ultra LLM, with Llama3-8B as the generator. Synthetic Q&A pairs were made with NVIDIA TrueQuery; the 253B Nemotron was used for TXT2KG and ModernBERT for vector embeddings. Approximate bar heights read from the plot:

• Basic RAG ≈ 70; Basic RAG + fine-tuned LLM ≈ 77; Agentic RAG (NV-AIQ) ≈ 73;
• Basic Graph RAG ≈ 59 (the lowest); Basic Graph RAG + fine-tuned LLM ≈ 77;
• GNN+LLM RAG (also fine-tuned) ≈ 81–82, the best real system;
• "Golden retrieval" ≈ 82–83: the upper bound ("speed of light"), where the generator is fed the same documents that were used to create each Q&A pair.

Lessons: naive GraphRAG (just textify a subgraph) can be worse than plain RAG; fine-tuning helps everything; and the GNN+LLM system comes close to perfect retrieval.

Slide 39: a "vibe-coded" GUI. To show the codebase is easy to build on, the PM (Santosh Bhavani) built a GUI with Cursor: "Text-to-Graph on NVIDIA DGX Spark" (build.nvidia.com/spark/txt2kg). The screenshot shows a dark interface with a large rendered knowledge graph.

Slide 40: next steps for PyG's GNN+LLM work:

• keep improving the codebase (speed, usability, accuracy, understandability);
• accelerate retrieval with cuGraph + cuVS (GPU graph analytics and GPU vector search), to scale to a trillion edges and beyond;
• align TXT2KG with KGGen (a Stanford tool for extracting KGs from text);
• more modalities, new problems and new retrieval techniques.

Slides 41 and 44: the general idea. Using GNN embeddings as a prefix to a Transformer/LLM prompt is very general. Examples:

1. graphs = molecules or cells; language task = biology, chemistry, drug discovery;
2. graphs = customer data; task = "talk to your customer data";
3. graphs = documents with hyperlinks; task = "talk to your docs".

Graphs can also mix modalities: Amazon products where each node has a text review and a photo; relational-database heterographs as in RelBench and Kumo.ai (Lectures 13–14). Node features could be text (natural language or code), images, audio, video, or molecule/cell embeddings.

Slide 42: a scientific GNN+LLM community sprint (biology and chemistry): the goal is to add GNN+LLM support for the sciences, with 3 biology papers and 1 chemistry paper reimplemented, to advance medicine and science (see examples/llm/README.md).

Slide 43: MoleculeGPT, "talk to your molecule". A user asks "How is the color and odor of molecule C1-CC-C(C-C1)C=O?" (a SMILES string) and the bot answers "It appears a clear colorless to yellow liquid with a bitter almond odor." A GNN encodes the molecular graph; its embedding is fed to the LLM through a Q-Former (the bridge module from BLIP-2; QR codes link to the PyG example, the MoleculeGPT paper and the Q-Former paper).

Slide 45 draws a spectrum. On the left: unstructured documents turned into a KG by an LLM (TXT2KG): less accurate, because extraction makes mistakes. On the right: natively structured graph data such as relational databases (RelBench), hyperlinked web pages and molecules: highly accurate. In between: the two kinds likely cover different knowledge, and most enterprises have both. The proposal is to combine them, with a separate GNN per graph type.

Slide 46: graph transformers vs message passing (future). Today there are two kinds of GNN: traditional message passing and the newer graph transformers (with a QR code to a webinar). More graph transformers are coming to PyG; a future goal is to measure how they change the accuracy of the GNN+LLM systems above.

Slide 47: conclusion.

• General AI rule: data is king. Example: even the best model cannot learn from randomly labelled data.
• Key takeaway: structured data (GNN) + unstructured data (LLM) = better GenAI.
• All of this is built and optimized in NVIDIA's PyG container (catalog.ngc.nvidia.com).