An LLM-Guided Query-Aware Inference System for GNN Models on Large Knowledge Graphs

This paper introduces KG-WISE, an LLM-guided inference system that accelerates Graph Neural Network queries on large Knowledge Graphs by dynamically decomposing models into fine-grained components and retrieving only semantically relevant subgraphs, achieving up to 28x faster inference and 98% lower memory usage while maintaining accuracy.

The Gist

Imagine you have a massive, ancient library containing every fact about the world (a Knowledge Graph). Inside this library, you have a brilliant librarian (a Graph Neural Network or GNN) who is incredibly smart at finding connections between things.

However, there's a problem: The library is so huge (millions of books and shelves) that every time you ask the librarian a simple question like, "Who wrote this specific book?", they have to:

1. Walk to the front desk.
2. Drag the entire library out of the building and onto a table.
3. Read every single book to find the answer.
4. Put the whole library back.

This is slow, exhausting, and wastes a massive amount of energy. Even if you only need to know about one author, the librarian still moves the whole building.

Enter KG-WISE.

The paper introduces a new system called KG-WISE that changes how we ask questions and how the librarian works. Here is how it works, using simple analogies:

1. The "Smart Assistant" (The LLM)

Instead of the librarian guessing which books you might need, KG-WISE uses a Super-Smart Assistant (a Large Language Model, or LLM) to help you.

• Before: You ask, "Who wrote this?" and the librarian grabs everything.
• Now: You tell the Assistant, "I need info on Author X." The Assistant instantly understands the context. It knows exactly which shelves contain Author X's books and which ones are irrelevant (like books about cooking or space travel).
• The Magic: The Assistant writes a special "shopping list" (a Query Template) that tells the librarian exactly which few books to pull out. This list is created once and reused, so it's very fast.

2. The "Modular Library" (Model Decomposition)

In the old system, the librarian's knowledge was stored in one giant, heavy backpack. You had to carry the whole backpack to find one fact.

KG-WISE breaks the backpack apart.

• It separates the rules (how to read) from the facts (the specific book contents).
• It stores these facts in a high-tech, organized warehouse (a Key-Value Store) where every single book has its own labeled slot.
• Now, the librarian doesn't need the whole backpack. They just need to grab the specific books on the "shopping list."

3. The "On-Demand Delivery" (Query-Aware Inference)

When you ask a question:

1. The system looks at your "shopping list" (the query template).
2. It goes to the warehouse and only loads the specific books (data) and the specific pages of the rulebook (model weights) needed for that question.
3. The librarian reads just those few pages, gives you the answer, and puts everything back.

The Result:

• Speed: Instead of moving the whole library, they only moved a few boxes. The paper says this is 28 times faster.
• Memory: Instead of needing a warehouse to hold the whole library, they only need a small shelf. This uses 98% less memory.
• Accuracy: Because they aren't distracted by thousands of irrelevant books, the librarian actually gives better answers in some cases.

Why is this a big deal?

Imagine trying to find a needle in a haystack.

• Old Way: You burn the whole haystack to find the needle. (Expensive, slow, wasteful).
• KG-WISE Way: You use a metal detector to find exactly where the needle is, dig only that spot, and leave the rest of the hay untouched.

The Bottom Line

KG-WISE is like upgrading from a delivery truck that carries your entire house to your new apartment, to a drone that delivers only the specific box you ordered. It makes using super-smart AI on massive data sets fast, cheap, and energy-efficient, without losing any of the intelligence.

In short: It stops the computer from doing unnecessary work by asking, "Do you really need all this data, or just this part?" and then only loading the part you need.

Technical Summary

Here is a detailed technical summary of the paper "An LLM-Guided Query-Aware Inference System for GNN Models on Large Knowledge Graphs" (KG-WISE).

1. Problem Statement

Graph Neural Networks (GNNs) are critical for applications on Knowledge Graphs (KGs), such as recommendation systems and drug discovery. However, performing inference on large-scale heterogeneous KGs faces significant scalability challenges:

• Monolithic Loading: Existing systems store trained GNNs as single monolithic files. Even for a query involving a small set of target nodes, the system must load the entire model (including weights and precomputed embeddings for all non-target nodes) and the full graph adjacency into memory.
• Redundant Computation: In heterogeneous KGs, precomputed embeddings often constitute >95% of the model size. Existing acceleration techniques (pruning, quantization, knowledge distillation) reduce model weights but fail to address the cost of loading irrelevant embeddings for non-target nodes.
• Lack of Adaptivity: Current methods treat all queries uniformly (e.g., fixed $L$-hop neighborhoods), ignoring the semantic structure and density variations specific to each query. This leads to excessive data loading and redundant computation on large, heterogeneous graphs.

2. Methodology: KG-WISE

KG-WISE introduces a task-driven, query-aware inference paradigm that adapts computation and data loading to the specific structure and semantics of each inference query. The system operates in three main phases:

A. LLM-Guided Query Template Generation (Training Phase)

Instead of training on the full KG or using random sampling, KG-WISE uses a Large Language Model (LLM) to generate a reusable SPARQL query template ($QT$) for a specific task.

• Process: The LLM analyzes the task description and KG schema statistics to identify semantically relevant entities and relations.
• Output: It generates a constrained SPARQL template that extracts a compact, dense subgraph relevant to the task.
• Benefit: This subgraph is significantly smaller than the full KG, reducing training time and memory. The template is stored for reuse during inference, eliminating the need for LLM calls at runtime.

B. Fine-Grained Model Decomposition & Storage

KG-WISE decouples the trained GNN model into modular components and stores them separately to enable partial loading:

• Decomposition: The model is split into Model Parameters (weights/biases) and Node Embeddings.
• Storage Architecture:
  • RDF Engine: Stores the KG schema, metadata, and the reusable SPARQL query templates.
  • Key-Value (KV) Store (Zarr): Stores node embeddings. Crucially, embeddings are chunked by node type and indexed by node ID.
  • Parameter Store: Stores model weights separately.
• Result: This allows the system to retrieve only the specific embedding chunks required for a query, rather than loading the entire embedding matrix.

C. Query-Aware Inference Pipeline

When an inference query arrives (targeting a set of nodes $TN$):

1. Subgraph Extraction: The system retrieves the stored SPARQL template, instantiates it with the target nodes, and executes it against the RDF engine to fetch the semantic subgraph ($SG$).
2. Compact Model Instantiation:
   • The system loads only the model weights and the specific embedding chunks corresponding to nodes in $SG$.
   • It constructs a sparse tensor for embeddings, where only relevant nodes have values, and irrelevant nodes are implicitly omitted.
3. Adaptive Aggregation: Based on the sparsity of the retrieved subgraph, the system dynamically selects between sparse or dense tensor aggregation to optimize computation.
4. Prediction: The forward pass is executed on this compact, query-specific model ($\hat{M}$), producing predictions with minimal memory footprint.

3. Key Contributions

1. First End-to-End Query-Aware System: The first system designed specifically for scalable GNN storage and inference on large heterogeneous KGs that tailors data loading to query semantics.
2. Fine-Grained Decomposition: A novel storage mechanism that decouples embeddings from weights, enabling partial model loading with minimal overhead.
3. LLM-Guided Subgraph Extraction: Utilizes LLMs to generate reusable query templates that extract semantically relevant subgraphs, ensuring the model is instantiated only with necessary components.
4. Scalable Architecture: Combines an RDF engine for semantic retrieval with a KV-store for efficient numerical tensor access, supporting both CPU and GPU inference.

4. Experimental Results

The authors evaluated KG-WISE on six large real-world KGs (up to 42 million nodes and 166 million edges) across Node Classification (NC) and Link Prediction (LP) tasks.

• Performance Gains:
  • Inference Speed: Up to 28× faster than state-of-the-art (SOTA) systems (e.g., GraphSAINT, GCNP, Degree-Quant).
  • Memory Usage: Up to 98% reduction in memory consumption. For example, on the DBLP dataset, model instantiation dropped from ~6.9 GB (baselines) to as low as 17 MB for small queries.
  • Accuracy: Maintained or improved accuracy compared to SOTA methods. In some cases, the sparsity of the extracted subgraph acted as regularization, improving accuracy.
• Scalability:
  • Weak Scalability: Demonstrated better memory efficiency (0.26 vs 0.24) compared to GraphSAINT.
  • Strong Scalability: Showed higher computational speedup (2.4× vs 1.83×) when increasing worker cores.
• LLM Agnosticism: The system works effectively with both commercial (Gemini, GPT-4/5) and open-weight (Qwen, DeepSeek) LLMs. While proprietary models generated slightly larger subgraphs, open-weight models produced more compact subgraphs with comparable accuracy, proving the system is not tied to a specific LLM.
• Sustainability: KG-WISE reduced energy consumption by 62% and CO2 emissions by 60% compared to baseline methods.

5. Significance

KG-WISE addresses a fundamental bottleneck in deploying GNNs on large-scale industrial knowledge graphs. By shifting from "load everything" to "load only what is needed," it enables:

• Cost-Effective Deployment: Drastically reduces hardware requirements (RAM/VRAM), allowing large KG inference on consumer-grade or smaller server hardware.
• Real-Time Responsiveness: The massive reduction in latency makes real-time GNN inference feasible for dynamic applications.
• Sustainability: Significant energy savings contribute to greener AI practices.
• Paradigm Shift: It moves the field from static model compression (pruning/quantization) to dynamic, semantic-aware model instantiation, proving that LLMs can effectively guide the structural optimization of traditional ML pipelines.