SAKE: Structured Agentic Knowledge Extrapolation for Complex LLM Reasoning via Reinforcement Learning

1. Problem Definition

The paper addresses the challenge of Knowledge Extrapolation: the ability to infer novel information by combining and extending existing, explicitly available knowledge. This is critical in specialized domains (e.g., biomedicine) where retrieving a complete external knowledge base is impractical.

• The Gap: Current Large Language Models (LLMs) often fail in knowledge-intensive tasks because:
  • Parametric Knowledge: Pre-trained weights are insufficient for niche facts.
  • Retrieval Limitations: Standard Retrieval-Augmented Generation (RAG) assumes the answer exists in the retrieved context. If the Knowledge Graph (KG) is incomplete, the model cannot answer.
  • Complexity: Existing agentic frameworks (e.g., GraphRAG, ToG) rely on complex, multi-step prompting that is difficult for smaller models (3B–7B parameters) to follow reliably, leading to hallucinations or high token costs.
• Goal: To enable small, open-weight models to autonomously retrieve structured knowledge, identify structural patterns, and extrapolate new facts via analogy to answer questions where the direct answer is missing from the KG.

2. Methodology: SAKE Framework

SAKE (Structured Agentic Knowledge Extrapolation) is a tool-augmented Reinforcement Learning (RL) framework. It trains LLMs to perform a multi-turn reasoning pipeline without Supervised Fine-Tuning (SFT).

Core Components

1. Tool-Augmented Environment:
   The model interacts with a Knowledge Graph (KG) via two deterministic external tools:

   • Tool 1 (Entity Group Construction): Links extracted query entities to the KG and retrieves semantically similar "sibling" concepts (forming entity groups). This bridges the gap between query terminology and KG vocabulary.
   • Tool 2 (Cross-Group Triplet Retrieval): Retrieves KG triplets that connect entities across different selected groups, providing structural evidence for reasoning.

2. Three-Turn Agentic Pipeline:
   The model generates a trajectory interleaved with tool calls:

   • Turn 1 (Extraction): The model extracts key entities from the query (e.g., specific diseases or molecules) and triggers Tool 1.
   • Turn 2 (Filtering): The model receives entity groups, evaluates their relevance, filters out noisy groups, and triggers Tool 2 to retrieve cross-group triplets.
   • Turn 3 (Associative Reasoning): The model analyzes retrieved triplets. If the answer is not directly present, it performs knowledge extrapolation by substituting semantically similar concepts from the entity groups into existing triplet patterns (e.g., mapping (hormone, treats, mental_disorder) to (melatonin, treats, insomnia)).

3. Training Objective (GRPO):

   • Algorithm: The framework uses Group Relative Policy Optimization (GRPO) within the VERL framework.
   • End-to-End Optimization: The entire trajectory (extraction, filtering, reasoning) is optimized jointly. Tool outputs are deterministic and masked from the loss, but the final outcome reward supervises upstream decisions.
   • Curriculum Reward: A three-phase reward function guides training:
     1. Format Phase: Rewards correct structural tagging (ensuring the pipeline runs).
     2. Transition Phase: Rewards both format and answer accuracy.
     3. Accuracy Phase: Rewards only answer correctness (format compliance becomes habitual).

3. Key Contributions

• Learnable Knowledge Extrapolation: Demonstrates that the ability to extrapolate missing knowledge via analogy is a behavior that can be learned end-to-end by small models (3B/7B) through RL, rather than requiring complex prompt engineering or massive parameter counts.
• Unified Agentic Policy: Unifies retrieval and reasoning into a single policy. Unlike prior works that separate retrieval and reasoning or rely on pre-computed data, SAKE learns what to retrieve and how to reason over it dynamically.
• Efficiency: Achieves state-of-the-art performance with a single inference pass, drastically reducing token usage compared to iterative agentic frameworks.
• No SFT Required: The model is trained directly from a base instruct model using only RL, proving that multi-turn agentic behavior can emerge purely from reward-driven exploration.

4. Experimental Results

The authors evaluated SAKE on two distinct Knowledge Graphs:

• Biomedical: UMLS (135 nodes, 5.9K triplets) on PubMedQA, BioASQ, and ProcessBank.
• Commonsense: ConceptNet (844K nodes, 2.1M triplets) on CommonsenseQA.

Key Findings:

• Performance: A Qwen2.5-7B model with SAKE outperformed GPT-3.5-Turbo paired with state-of-the-art agentic frameworks (GIVE):
  • Biomedical: 75.4% vs. 70.1% (Weighted Average).
  • Commonsense: 81.3% vs. 74.7%.
• Scalability: SAKE improved performance consistently across both 3B and 7B models, whereas other baselines (like SFT or standard GRPO) showed unstable gains or degradation at different scales.
• Efficiency: SAKE reduced token usage by >90% compared to iterative frameworks like ToG and GIVE (e.g., ~800 tokens vs. ~14,000 tokens per question).
• Ablation Studies:
  • Group Filtering: Removing this step reduced accuracy by ~2.5–5.0%, confirming the need to filter noise before retrieval.
  • Agentic Retrieval: Learned retrieval outperformed pre-computed retrieval by ~5–7%, proving the value of query-adaptive extraction.
  • Knowledge Extrapolation: Removing the extrapolation step caused the largest drop (up to 15.4% on BioASQ), validating that the ability to construct new triplets is essential for incomplete KGs.

5. Significance

SAKE represents a paradigm shift in deploying LLMs for specialized domains:

1. Democratization of Reasoning: It proves that small, open-weight models can achieve reasoning capabilities previously thought to require large models or complex, multi-step prompting strategies.
2. Handling Incompleteness: It offers a robust solution for domains where knowledge bases are inherently incomplete, moving beyond simple retrieval to associative reasoning.
3. Cost-Effectiveness: By reducing token usage by over 90% and requiring only one inference pass, SAKE makes advanced agentic reasoning economically viable for real-world applications.

The code is available at https://anonymous.4open.science/r/SAKE-7585.