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

Waleed Afandi; Hussein Abdallah; Ashraf Aboulnaga; Essam Mansour

cs.LGMar 4, 2026

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

Authors: Waleed Afandi, Hussein Abdallah, Ashraf Aboulnaga, Essam Mansour

Organizations: Concordia University · University of Texas at Arlington

Abstract

Efficient inference for graph neural networks (GNNs) on large knowledge graphs (KGs) is essential for many real-world applications. GNN inference queries are computationally expensive and vary in complexity, as each involves a different number of target nodes linked to subgraphs of diverse densities and structures. Existing acceleration methods, such as pruning, quantization, and knowledge distillation, instantiate smaller models but do not adapt them to the structure or semantics of individual queries. They also store models as monolithic files that must be fully loaded, and miss the opportunity to retrieve only the neighboring nodes and corresponding model components that are semantically relevant to the target nodes. These limitations lead to excessive data loading and redundant computation on large KGs. This paper presents KG-WISE, a task-driven inference paradigm for large KGs. KG-WISE decomposes trained GNN models into fine-grained components that can be partially loaded based on the structure of the queried subgraph. It employs large language models (LLMs) to generate reusable query templates that extract semantically relevant subgraphs for each task, enabling query-aware and compact model instantiation. We evaluate KG-WISE on six large KGs with up to 42 million nodes and 166 million edges. KG-WISE achieves up to 28x faster inference and 98% lower memory usage than state-of-the-art systems while maintaining or improving accuracy across both commercial and open-weight LLMs.

Figures & tables

Fig. 1: A GNN inference query for target nodes ( $TN$ ) loads the KG’s adjacency, model, and embedding matrices from storage, then performs message passing and embedding updates to produce predictions. This process is resource-intensive and scales poorly with the size of these matrices in KGs.

	Pruning (GCNP) [ 8 ]	Quantization (DQ) [ 9 ]	Know. Distillation (GKD) [ 10 ]	Ours (KG-WISE)
Optimization Overhead	Per layer Channel Pruning	Simulation of Quantization effect	Student Model Training	Decompose $H$ into KV Store
Additional Training	$\bm{\checkmark}$	$\bm{\checkmark}$	$\bm{\checkmark}$	$\bm{\times}$
Storage Mechanism	Disk based single file	Disk based single file	Disk based single file	Disk based & Key-Value store
Partial Graph Loading	$\bm{\times}$ (Full Graph)	$\bm{\times}$ (Full Graph)	$\bm{\checkmark}$ (Randomly sparsed)	$\bm{\checkmark}$ (Query based)
Partial Embedding Loading	$\bm{\times}$ (Full Graph)	$\bm{\times}$ (Full Graph)	$\bm{\times}$ (Full Graph)	$\bm{\checkmark}$ (1 hop)
Forward Pass	Full	Full	Random Sparse	Query Aware

TABLE I: A Comparative Analysis of Existing and Proposed Inference Acceleration Techniques.

Fig. 2: KG-WISE orchestrates training and inference on large KGs through LLM-guided subgraph extraction, fine-grained model storage, and query-aware model instantiation.

Fig. 3: The top half shows an encoded KG and a GNN trained on it. The GNN’s main components are the learned parameters and non-target node embeddings . The bottom half illustrates how KG-WISE maps the encoding and decouples the GNN components. The node embeddings are stored in a key-value store as as row-wise chunks grouped by node type in a Zarr KV store.

Fig. 4: KG-WISE inference pipeline. Given an inference query, KG-WISE loads a stored SPARQL template to extract a semantically relevant subgraph $SG$ , then instantiates a compact model $\widetilde{M}$ by loading only the required embeddings and weights from the KV store. Inference is executed on-demand using sparse tensor aggregation over $SG$ , avoiding full model loading.

TT	Name	KG	#nodes	#edges	#n-type	#e-type	Metric
NC	PV	MAG-42M	42.4M	166M	58	62	ACC
NC	PC	YAGO4	30.7M	400M	104	98	ACC
NC	PV	DBLP-15M	15.6M	252M	42	48	ACC
LP	AA	DBLP-15M	15.6M	252M	42	48	H@10
LP	PO	ogbl-wikikg2	2.5M	17M	9.3K	535	H@10
LP	CA	YAGO3-10	123K	1.1M	23	37	H@10

TABLE II: KGs statistics and GNN tasks: Task Types (TT) include node classification (NC) and link prediction (LP).

Task	Model statistics	Baseline Accelerators				KG-WISE’s model with different — $TN$ —
Task	Model statistics	GS	IBMB	GCNP	DQ	100	200	400	800	1600
DBLP (PV) NC	# Parameters	1.86 B	2 B	1.85 B	1.86 B	430.42 M	430.42 M	430.42 M	430.42 M	430.42 M
	Size in disk	6.9 GB	7.4 GB	6.9 GB	6.9 GB	17.1 MB	34.8 MB	50.5 MB	67 MB	121.1 MB
	Size(Weights)/Size(Model)	0.01	0.01	0.01	-	0.5	0.33	0.2	0.11	0.06
	Size(Embeddings)/Size(Model)	0.99	0.99	0.99	-	0.5	0.66	0.79	0.88	0.93
YAGO4 (PC) NC	# Parameters	3.65 B	B	1.83 B	3.65 B	243.7 M	243.7 M	243.7 M	243.7 M	243.7 M

TABLE III: Model size breakdown for KG-WISE and baseline accelerators. In KG-based GNNs, most of the model size comes from non-target node embeddings, which baselines always load in full (constant model size across all

|TN|

). In contrast, KG-WISE loads only embeddings for the extracted subgraph, yielding a query-dependent model size. “OOM” denotes out-of-memory and “N/A” not applicable. For DQ, weights and embeddings cannot be decoupled due to shared quantization statistics.

Fig. 5: Performance across NC tasks is based on three metrics: (A) Inference Accuracy (higher is better), (B) Inference-Time (lower is better), and (C) Inference Memory (lower is better). The top and middle sections illustrate the results for the Paper-Venue task on DBLP and MAG, respectively. The bottom figures present the results of the Place-Country task on YAGO4. The inference query performs inference for 1K target nodes stratified across all classes. KG-WISE archives comparable inference accuracy compared with the SOTA training/inference accelerators (Graph SAINT, IBMB, GCNP, and DQ). KG-WISE outperforms the SOTA methods by up to 28x in inference time on YAGO4 with memory reduction up to 98%.

Panel of Figure 9 — Fig. 6: Results on link prediction tasks across three KGs, reporting (A) Hits@10, (B) inference time, and (C) memory usage. Rows correspond to YAGO3-10 (CA), DBLP (AA), and WikiKG (PO). Each query uses 100 target nodes. KG-WISE improves accuracy while reducing inference time (up to 7 $\times$ ) and memory compared to MorsE and DQ (which fails OOM in two cases).

Fig. 7: The inference time and memory consumption of KG-WISE and Graph-SAINT inference pipelines under weak and strong scalability settings. KG-WISE shows better scalability in both settings.

Panel of Figure 11 — Fig. 8: Evaluation of KG-WISE while running on a CPU and a GPU

Panel of Figure 12 — Fig. 9: KG-TOSA vs. KG-WISE on node classification tasks using three metrics: (A) Inference Accuracy (higher is better), (B) Inference Time (lower is better), and (C) Inference Memory (lower is better). Our LLM-guided subgraph extraction and query-aware inference substantially improve efficiency: KG-WISE reduces inference time by 41–52% and memory usage by 60–84% while preserving accuracy.

Fig. 10: The number of KV-store chunks loaded per node type by KG-WISE to answer the 1K target node inference query on the DBLP KG. KG-WISE loaded just 10% of the KV-chunks, achieving significant savings in both loading time and memory.

Panel of Figure 14 — Fig. 11: Evaluation of different LLMs used to generate the query template $Q_{T}$ in KG-WISE across three KG tasks. Accuracy remains stable across all models, while inference time and memory vary due to differences in predicate selection. Open-weight LLMs (Qwen, GPT-oss) produce more compact subgraphs on heterogeneous KGs (e.g., YAGO4), yielding lower memory usage and faster inference compared to proprietary LLMs.

Metric	GraphSAINT	KG-WISE
Total Energy (Wh)	1.19	0.42
Energy per Target Node (Wh)	1.2 $\times 10^{-3}$	0.4 $\times 10^{-3}$
Total Emission (g CO 2 )	1.5 $\times 10^{-1}$	0.53 $\times 10^{-1}$
Emission per Target Node (g CO 2 )	1.5 $\times 10^{-4}$	0.53 $\times 10^{-4}$
Relative CO 2 Efficiency	-	2.83x

TABLE IV: CO 2 emissions and energy consumption of KG-WISE and Graph-SAINT. KG-WISE generates 60% less CO 2 emissions and consumes 62% less energy.

Fig. 14: Training metrics on NC tasks is measured using three metrics: (A) Test Accuracy (higher is better), (B) Training Time (lower is better), and (C) Max Memory Usage (lower is better). The top and middle sections show results for the Paper-Venue task on DBLP and MAG, respectively, while the bottom presents results for the Place-Country task on YAGO4.KG-WISE achieves comparable accuracy to SOTA methods and outperforms them with up to 6x faster training and 90% memory reduction on YAGO4.

Explore similar work

Search-on-Graph-R1: Training Large Language Models to Search Knowledge Graphs with Reinforcement Learning
Jul 20, 2026Jia Ao Sun, Hao Yu, Fengran Mo +4LLM Reasoning with Graphs RL for Language Model Reasoning
Université de Montréal · Mila – Québec AI Institute · McGill University +2
Code-on-Graph: Iterative Programmatic Reasoning via Large Language Models on Knowledge Graphs
Jun 2, 2026Weiwei Ding, Zixuan Li, Long Bai +7LLM Grounding LLM Reasoning with Graphs
1Key Laboratory of AI Safety, Institute of Computing Technology, Chinese Academy of Sciences · 2Shandong University · 3Shandong University-Weihai Research Institute of Industrial Technology
Leveraging LLM-GNN Integration for Open-World Question Answering over Knowledge Graphs
Apr 15, 2026Hussein Abdallah, Ibrahim Abdelaziz, Panos Kalnis +1Knowledge Graph Question Answering LLM Reasoning with Graphs
Concordia University · IBM · KAUST

Jul 20, 2026J/K move · Enter open · S save

Search-on-Graph-R1: Training Large Language Models to Search Knowledge Graphs with Reinforcement Learning

Jia Ao Sun, Hao Yu, Fengran Mo, Zhan Su+3

Université de Montréal · Mila – Québec AI Institute · McGill University+2

LLM Reasoning with Graphs RL for Language Model Reasoning Graph-Based Retrieval LLM Tool Use Knowledge Graph Question Answering

Open paper AskPDF arXiv

Knowledge graph question answering (KGQA) requires navigating from topic entities to an answer several relations away. Recent methods prompt a frontier LLM to explore the graph through a retrieval tool, but their reliance on frontier-scale inference makes them costly to deploy. We present Search-on-Graph-R1 (\sogrone{}), which internalizes this navigation into a compact 8B model through supervised fine-tuning (SFT) followed by reinforcement learning (RL). Our central idea is to scaffold a frontier teacher with each question's gold SPARQL query, so the teacher traverses a known answer-bearing path with a live Search tool rather than having to discover the path itself. Since every call executes against a live Freebase server, the resulting trajectories are grounded in the knowledge graph by construction. On WebQSP, CWQ, and GrailQA, \sogrone{} at 8B surpasses every frozen frontier-LLM system in our comparison and posts the strongest results on CWQ of any system we compare against. It does so using no auxiliary module at inference and no LLM judge during training. Isolating each training stage shows that SFT and RL contribute complementary gains, our approach transfers across model families, and RL learns to reach answers in fewer Search calls than its SFT initialization.

Abstract

Explore similar work

Search-on-Graph-R1: Training Large Language Models to Search Knowledge Graphs with Reinforcement Learning

Code-on-Graph: Iterative Programmatic Reasoning via Large Language Models on Knowledge Graphs

Leveraging LLM-GNN Integration for Open-World Question Answering over Knowledge Graphs