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

TL;DR

This paper addresses hallucinations in knowledge graph question answering by focusing on small LLMs and proposing RPO-RAG, a KG-based retrieval augmented generation framework. It introduces three innovations: a query-path semantic sampling strategy to align training paths with query intent, a relation-aware weighted preference optimization to supervise intermediate reasoning steps, and an answer-centered prompt design that groups evidence by candidate answers. Experiments on WebQSP and CWQ show that RPO-RAG consistently improves accuracy and reduces the gap to GPT-based models, with strong gains for models as small as 1B–3B parameters and favorable efficiency trade-offs. The work highlights a practical path to resource-efficient KGQA, supported by dataset quality analyses that confirm semantically guided supervision yields better retriever signaling and more faithful reasoning.

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.

Figures (8)

• Figure 1: An example of the prompt used in existing works (top) and our designed prompt (bottom), each shown with a predicted answer from a small LLM (Llama3.2-3B ). Paths (B1)–(B16) correspond to reasoning paths in existing works, while (O1)–(O4) denote our answer-centered reasoning paths. Angle brackets mark candidate answers (e.g., <Bishop> and <Q>); (O2) and (O3) group paths under these candidates, respectively.
• Figure 2: Overview of the RPO-RAG framework. (1) Query-Path Semantic Sampling: constructs query-aligned training paths via dynamic clustering to capture query intent. (2) Semantic-Matching Retriever: retrieves reasoning paths semantically consistent with the query using a pretrained language model. (3) Dual-Objective Optimization: optimizes relation-level preference and answer-centered prompt objectives to align small LLMs with structured reasoning.
• Figure 3: An example of Query-Path Semantic Sampling.
• Figure 4: Illustration of Relation-aware Weighted Preference Optimization. For the same question and current path (left), candidate next relations (right) are scored by semantic relevance. Higher-scored relations are treated as preferred (CHOSEN), while semantically misaligned ones are treated as non-preferred (REJECT).
• Figure 5: Comparison of query–path semantic alignment between datasets under top-1 and top-3 settings.