RPO-RAG: Aligning Small LLMs with Relation-aware Preference Optimization for Knowledge Graph Question Answering

AI summary

Overview

• Research area: Knowledge Graph Question Answering (KGQA), retrieval-augmented generation (RAG), preference optimization, and efficient small language models.
• Technical level: Advanced (assumes familiarity with knowledge graphs, RAG pipelines, preference optimization such as DPO/SimPO, and chain-of-thought style reasoning over retrieved evidence).
• Scope: The paper proposes RPO-RAG, a KG-based RAG framework that tailors retrieval sampling, training objectives, and prompt layout specifically to sub-8B LLMs, and evaluates it on the WebQSP and CWQ benchmarks grounded in Freebase.

What This Paper Is About

Large language models still hallucinate on knowledge-intensive questions, and KG-based RAG is a common fix, but existing methods pick training paths with semantics-unaware heuristics, supervise only the final answer rather than intermediate reasoning, and dump retrieved paths into the reasoner as a flat ungrouped list. Those weaknesses hurt small LLMs most, and prior work mostly assumes backbones above 7B parameters or large proprietary models like ChatGPT/GPT-4. RPO-RAG's goal is to make sub-8B models reason over knowledge graphs accurately by aligning the entire retrieval-to-reasoning pipeline with the structure of the graph.

Key Contributions

1. RPO-RAG framework for small LLMs: The authors describe it as the first KG-based RAG framework specifically designed for small LLMs, introducing a query-path semantic sampling strategy that automatically measures semantic relevance between queries and reasoning paths to build higher-quality training data.
2. Relation-level preference optimization: RPO-RAG explicitly models the LLM's reasoning process at the relation level, which the authors state is the first framework to incorporate relations into preference optimization for KG-based RAG.
3. Answer-centered prompt design: Retrieved entities and supporting paths are regrouped into answer-centered reasoning paths so that paths supporting the same candidate answer appear together, rather than as an unordered list.
4. Benchmark gains for compact models: Experiments on WebQSP and CWQ show consistent improvements over graph-based and small-LLM baselines and a narrowed gap to GPT-based systems, with gains even at 1B–3B parameters.

Main Findings

• WebQSP results: RPO-RAG (Llama3.1-8B) reaches 89.9 Hit and 81.3 F1, which the paper calls state-of-the-art among models up to 8B parameters, surpassing the previous best (GCR) by +2.7% Hit and +10.2% F1. RPO-RAG (Llama2-7B) is second-best overall at 88.3 Hit and 77.8 F1.
• CWQ results: RPO-RAG (Llama2-7B) improves Hit by +1.3% and F1 by +4.9% over GNN-RAG (Llama2-7B), and RPO-RAG (Llama3.1-8B) reaches the highest Hit (72.3) and F1 (64.5) among all baselines at or under 8B parameters.
• Small-backbone gains: Compared with the vanilla Llama3.2-3B, RPO-RAG (Llama3.2-3B) improves Hit by +24.8% on WebQSP and +46.1% on CWQ; RPO-RAG (Llama3.2-1B) yields +32.1% and +46% Hit gains over its vanilla counterpart.
• Even 1B models beat some GPT baselines: On WebQSP, RPO-RAG (Llama3.2-1B) achieves 82.3 Hit, surpassing ToG (ChatGPT) by +6.1%.
• Closing the GPT gap on CWQ: RPO-RAG (Llama3.1-8B) achieves 72.3 Hit and 64.5 F1 on CWQ, approaching GCR (ChatGPT) and reducing the gap to GCR (GPT-4o-mini) to within roughly 3–4 points.
• Efficiency on CWQ: Measured on a single NVIDIA RTX 3090 GPU, RPO-RAG (Llama2-7B) records Hit 68.1 with retrieval 0.07s, reasoning 1.36s and total 1.43s, versus GNN-RAG (Llama2-7B) at Hit 66.8 and 1.06s total, SubgraphRAG (Llama3.1-8B) at Hit 57.0 with 0.02s retrieval but 6.14s total, and GCR (Llama3.1-8B) at Hit 60.5 with 6.84s retrieval and 7.4s total. RPO-RAG (Llama3.1-8B) attains Hit 72.3 with 1.1s retrieval and 1.17s reasoning.
• Both components are necessary (ablation): For RPO-RAG (Llama2-7B), removing relation-aware optimization drops WebQSP to 81.6 Hit / 66.5 F1 and CWQ to 61.1 Hit / 54.2 F1; removing the answer-centered prompt drops WebQSP to 80.3 Hit / 65.7 F1 and CWQ to 58.6 Hit / 47.5 F1. The same pattern holds for RPO-RAG (Llama3.2-3B), where the two ablations give 78.8/63.4 and 78.1/62.7 on WebQSP and 58.1/55.9 and 57.3/46.6 on CWQ.
• Relation-aware optimization helps more on WebQSP; answer-centered prompts help more on CWQ: The paper attributes this to precise reasoning alignment versus multi-hop compositional reasoning respectively.
• Higher-quality training data: Versus the dataset released by RoG on CWQ, the proposed dataset keeps comparable recall while improving precision by +7.7% and F1 by +4.3% (RoG-cwq: precision 31.5, recall 35.2, F1 30.8; Ours-cwq: precision 39.2, recall 35.3, F1 35.1).
• Stronger semantic alignment: Using SBERT cosine similarity between query and candidate path embeddings, top-1 scores are comparable between datasets, but the gap widens to 0.04 at top-3, and the similarity drop from top-1 to top-3 is smaller for the proposed dataset (0.03 vs. 0.05).
• More accurate retrieval: The proposed retriever achieves 87.4% accuracy with an average of 116 retrieved paths, versus RoG at 69.6% with ARP 27 and GNN-RAG at 65.0% with ARP 20 — a +22.4% accuracy improvement over GNN-RAG. The paper argues the larger ARP reflects broader but semantically consistent coverage rather than redundancy.

Methodology in Plain English

RPO-RAG has three connected pieces.

1. Query-path semantic sampling. Instead of using shortest-path heuristics (e.g., BFS) to build training paths, the method enumerates shortest paths between the question's topic entity and the answer entity in the knowledge graph, embeds each query and path with a pretrained language model, and measures similarity with cosine similarity. A gradient-based dynamic clustering algorithm partitions paths by their similarity distribution and automatically picks the number of clusters at the point of maximum curvature. The cluster whose centroid embedding is most similar to the query is kept, and only paths in that cluster become training data. In the worked example — "Which languages are spoken at the location where the film 'Shutter' occurs?" — clustering identifies an inflection point at K=3 and retains only the semantically relevant connections.

2. Semantic-matching retriever. A pretrained language model (Sentence-BERT is the default) is trained with weak supervision from query-answer pairs to score relations along paths. At inference, the retriever expands paths with dynamic beam search whose expansion size is adjusted by thresholds based on gaps between similarity scores, pruning less relevant candidates. Entity type information from the KG schema is used as a filter: the model is trained to predict top-K answer entity types, and the top-5 predictions exclude paths whose terminal entity does not match.

3. Dual-objective optimization. The reasoner is trained with two objectives. Relation-aware weighted preference optimization constructs preferred relation sets from the representative cluster and non-preferred sets from alternative clusters, weights each relation by its distance to the cluster centroid (with a decay rate and a scaling factor), and optimizes a margin-based preference objective normalized by relation length — supervising the next relation given a partial path rather than only the final answer. Answer-centered prompt optimization maximizes the likelihood of the correct answer given a prompt in which paths are grouped by their end entity, so evidence converging on the same candidate is presented together.

Experimental setup. Evaluation uses WebQSP (mostly 1–2 hop questions) and CWQ (up to 4-hop reasoning), both grounded in Freebase. Reasoners are Llama2-7B, Llama3.1-8B, Llama3.2-3B, and Llama3.2-1B, fine-tuned with LoRA on 2× NVIDIA RTX 4090 GPUs. Metrics are Hit and F1. Baselines span graph-based methods (GraftNet, NSM, SR+NSM, UniKGQA), vanilla LLMs, and LLM+KG methods (ToG, RoG, GCR, SubgraphRAG, GNN-RAG). Code is available at https://github.com/KaeHyun/RPO-RAG (archived at https://doi.org/10.5281/zenodo.18322650), and trained models are archived at https://doi.org/10.5281/zenodo.18322931. Dataset statistics are listed in the paper's Table 7 in Appendix B.1, but the specific counts are not included in the truncated content provided.

Why This Matters

Impact on research. The paper argues that aligning retrieval and reasoning at the relation level, and presenting evidence in a form small models can use, is a practical route to scalable, resource-efficient KGQA. It shifts attention from large proprietary reasoners toward sub-8B open-weight models and provides a weakly supervised way to build preference pairs without manual annotation.

Real-world applications:

• On-device or edge question answering where a 1B–3B model must run under tight memory and compute budgets.
• Enterprise knowledge-base assistants that query structured corporate graphs for factual, auditable answers.
• Conversational recommender systems that need multi-hop lookups over structured user-item-knowledge graphs.
• Web-scale search and question answering over structured data such as Freebase-style graphs, where latency and grounding matter.

Industry relevance. The efficiency numbers matter commercially: RPO-RAG (Llama2-7B) holds total latency at 1.43s with Hit 68.1 on CWQ, while LLM-driven retrieval in GCR costs 6.84s of retrieval time alone, and SubgraphRAG's fast retrieval still yields 6.14s total. The framework also runs on commodity hardware (2× RTX 4090 for training, a single RTX 3090 for the efficiency measurements), which lowers the barrier for deployment of grounded QA without frontier-model inference costs.

Future Directions

• Attribute-aware reasoning: The authors note the current pipeline does not explicitly encode or propagate attribute-conditioned evidence, leading to higher error rates on queries that require reasoning over attribute values, with the effect amplified for backbones at or under 3B parameters.
• Schema- or property-guided sampling: Extending path construction to integrate attribute-aware retrieval and schema/property guidance.
• Supervising relation and attribute consistency together: Designing training signals that cover both relation-level and attribute-level consistency, rather than relations alone.
• Prompt or adapter mechanisms for attributes: Developing ways for small LLMs to represent and exploit attribute information more reliably, since added parameter-efficient components may be needed for compact backbones.

Target Audience

Researchers and engineers working on KGQA, retrieval-augmented generation, and preference optimization, especially those focused on deploying small open-weight models (1B–8B) under compute constraints; practitioners building grounded question answering over Freebase-style knowledge graphs; and readers interested in how training-signal design — sampling, relation-level supervision, and prompt structure — affects reasoning quality in compact LLMs.

Authors’ abstract

Large Language Models (LLMs) have recently demonstrated remarkable reasoning abilities, yet hallucinate on knowledge-intensive tasks. Retrieval-augmented generation (RAG) mitigates this issue by grounding answers in external sources, e.g., knowledge graphs (KGs). However, existing KG-based RAG approaches rely on semantics-unaware path sampling and are weakly aligned with KG reasoning objectives, which limits further accuracy gains. They also feed retrieved paths directly into the reasoner without organizing them into answer-centered reasoning paths, hindering small LLMs' ability to leverage the retrieved knowledge. Furthermore, prior works predominantly rely on large LLMs (e.g., ChatGPT/GPT-4) or assume backbones above 7B parameters, leaving sub-7B models underexplored. We address this gap with RPO-RAG, the first KG-based RAG framework specifically designed for small LLMs, to the best of our knowledge. RPO-RAG introduces three key innovations: (1) a query-path semantic sampling strategy that provides informative supervisory signals; (2) a relation-aware preference optimization that aligns training with intermediate KG reasoning signals (e.g., relation); and (3) an answer-centered prompt design that organizes entities and reasoning paths in an interpretable format. Extensive experiments on two benchmark Knowledge Graph Question Answering (KGQA) datasets, WebQSP and CWQ, demonstrate that RPO-RAG effectively bridges the performance gap between small and large language models. On WebQSP, it improves F1 by up to 8.8%, reflecting enhanced answer precision, while on CWQ it achieves new state-of-the-art results among models under 8B parameters in both Hit and F1. Overall, RPO-RAG substantially improves the reasoning capability of small LLMs, even under 3B parameters-highlighting their potential for resource-efficient and practical on-device KGQA applications.