Technical Summary: Search-on-Graph-R1 (SOG-R1)

Problem Statement

Knowledge Graph Question Answering (KGQA) requires navigating from topic entities to answers that may be several relations away. While recent approaches have demonstrated that frontier Large Language Models (LLMs) can achieve competitive accuracy by interacting with Knowledge Graphs (KGs) via retrieval tools, these methods rely on expensive, frozen frontier models (e.g., GPT-4) that are impractical for production due to latency, cost, and access constraints.

Existing attempts to post-train smaller, open-source models to replicate this capability face significant limitations:

1. Semantic Parsing Approaches: These train models to emit a complete logical form (e.g., SPARQL) in a single pass. They struggle when the required schema elements were unseen during training.
2. Disjoint Retrieval and Reasoning: Some methods (e.g., EoG) extract static subgraphs offline, bounding accuracy by the recall of the preprocessing step. Others (e.g., KG-Hopper) use live search but rely on external LLM judges to verify reasoning traces, introducing stochasticity and computational overhead.
3. Parametric Fallback: Many systems allow models to fall back on parametric memory when retrieval fails, meaning the facts cited are not guaranteed to exist in the KG.

The core challenge is to internalize KG navigation into a compact model (e.g., 8B parameters) that performs live retrieval and reasoning without auxiliary modules, external judges, or reliance on parametric memory.

Methodology: SOG-R1

SOG-R1 addresses these challenges through a two-stage post-training pipeline applied to a compact instruction-tuned student model (Llama-3.1-8B-Instruct). The approach fuses retrieval and reasoning into a single loop where the model decides what to fetch via live Search calls.

1. SPARQL-Scaffolded Cold-Start Trace Generation

To initialize the model's ability to navigate the KG, the authors generate high-quality training trajectories using a "teacher" strategy:

• Teacher Setup: A frontier-scale teacher LLM (Qwen3-235B) is provided with the question and its gold SPARQL query.
• Scaffolding: Instead of asking the teacher to discover the path, the gold SPARQL query acts as a blueprint. The teacher uses the same live Search tool the student will use to traverse the known answer-bearing path.
• Grounding: Every tool call executes against a live Freebase server. This ensures the resulting trajectories are grounded in the KG by construction, containing valid tool calls, observations, and the final answer.
• Filtering: Trajectories are retained only if the final answer exactly matches the gold answer and all calls execute without error.

2. Stage 1: Supervised Fine-Tuning (SFT)

The student model is fine-tuned using Low-Rank Adaptation (LoRA) on the collected teacher trajectories.

• The training data is serialized as multi-turn conversations where assistant turns are Search tool calls and intervening turns are the returned observations.
• The model learns to produce tool calls and answers based on the context, without being trained to predict the tool outputs themselves.
• This stage produces SOG-SFT, establishing a strong navigation prior.

3. Stage 2: Reinforcement Learning (RL)

The SFT model is further refined using Group Relative Policy Optimization (GRPO) with a deterministic outcome reward.

• Reward Function: The reward is the product of three factors:
  1. Format Gate ($g$): Ensures the model calls the tool at least once and terminates cleanly with a "Final answer:" string.
  2. Exact-Match ($r_{em}$): Returns 1 if the predicted answer set matches the gold answer, 0 otherwise.
  3. Efficiency Factor ($f$): A multiplicative penalty for excessive tool calls (turns), encouraging shorter paths among correct trajectories.
• Key Constraint: The reward loop contains no LLM judge and no gold-path supervision; it relies solely on the correctness of the final answer and the efficiency of the path.

Tool Innovation

The authors introduce a Batched 1-Hop Neighbour Retrieval tool. Unlike previous tools that required one call per entity, this tool accepts a list of entities in a single call. This prevents the model from dropping entities when expanding multiple nodes in the same hop, reducing inference cost and failure rates.

Key Results

The model was evaluated on three standard Freebase-based benchmarks: WebQSP, CWQ, and GrailQA.

• Performance vs. Frozen Frontier Models: The 8B SOG-R1 surpasses every frozen frontier-LLM system (including GPT-4 based methods like ToG, PoG, and LMP) on all three benchmarks.
• Performance vs. Post-Trained Systems:
  • SOG-R1 achieves the strongest CWQ accuracy (88.7%) of any system compared, outperforming other post-trained models like ChatKBQA (86.0%) and EoG (82.6%).
  • On WebQSP and GrailQA, it achieves results comparable to or slightly below the best post-trained systems (EoG), but with the distinct advantage of performing live retrieval rather than reasoning over a static, pre-extracted subgraph.
• Comparison with KG-Hopper: Against the most similar concurrent system (KG-Hopper, also Llama-3.1-8B), SOG-R1 shows significant gains (+14.8 to +34.6 points), attributed to the absence of external judges and the integration of retrieval and reasoning.

Ablation Studies and Analysis

• Stage Contributions:
  • SFT provides the largest performance jump, lifting the model from a zero-shot baseline to near-state-of-the-art levels, confirming the value of SPARQL-scaffolded traces.
  • RL provides complementary gains, particularly on the complex CWQ dataset, improving accuracy further and refining the policy.
  • RL-Only (SOG-R1-ZERO): Applying RL directly to the base model without SFT initialization performs significantly worse than the SFT-initialized model, indicating that the cold-start traces are necessary for effective RL in this domain.
• Path Compression: A unique finding is that the RL stage teaches the agent to navigate more efficiently. SOG-R1 reduces the average number of Search calls per question compared to its SFT initialization (e.g., -15% on WebQSP) while simultaneously increasing accuracy. This contrasts with web-search RL methods where agents tend to issue more calls.
• Data Scale: The approach is data-efficient. Performance plateaus early; only a few thousand annotated questions are needed to reach near-optimal accuracy. RL provides the most significant uplift when the cold-start data pool is small.
• Backbone Generality: The pipeline transfers successfully to a different model family (Qwen3-4B), maintaining the performance hierarchy (Base < SFT < RL).

Significance and Claims

The paper claims that SOG-R1 successfully internalizes the "Search-on-Graph" paradigm into a compact, single model. Its primary contributions are:

1. Internalization: It transforms KG traversal from an inference-time orchestration task into a learned behavior within an 8B model's weights, eliminating the need for auxiliary modules or external judges.
2. Grounded Cold-Start: By scaffolding a teacher with gold SPARQL queries and live execution, it generates training data that is inherently grounded in the KG, avoiding the hallucination risks of purely generative traces.
3. Efficiency: It demonstrates that RL can learn to reach answers in fewer steps than the SFT initialization, optimizing for both correctness and latency.
4. State-of-the-Art Efficiency: It achieves competitive or superior accuracy to much larger frozen models and other post-trained systems at a fraction of the inference cost.

Limitations: The authors note that the evaluation is restricted to Freebase and assumes the availability of gold topic entities (oracle entity linking) and gold SPARQL queries during training. The transferability to graphs without such annotations (e.g., Wikidata) or the ability to perform entity linking autonomously remains untested.