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

Probleemstelling

Het vermogen om nieuwe feiten af te leiden door bestaande kennis te combineren en uit te breiden (kennisextrapolatie) is essentieel voor het oplossen van complexe vragen in gespecialiseerde domeinen. Huidige Large Language Models (LLM's) kampen echter met twee hoofdproblemen in dergelijke scenario's:

1. Onvolledige kennis: Parametrische kennis (uit pre-training) of opgehaalde context (bij het testen) biedt vaak geen volledige redeneerketen, vooral in domeinen waar het ophalen van uitgebreide externe kennis onpraktisch is.
2. Hallucinaties en inefficiëntie: Het vertrouwen op alleen parametrische kennis of het gebruik van complexe, multi-staps prompting (zoals bij bestaande agentic frameworks) leidt vaak tot hallucinaties of is extreem kostbaar in termen van tokens en rekenkracht. Bestaande methoden zoals GraphRAG of ToG gaan vaak uit van complete kennisgrafieken of vereisen dure, iteratieve prompts die kleine modellen niet betrouwbaar kunnen volgen.

Methodologie: SAKE Framework

De auteurs stellen SAKE (Structured Agentic Knowledge Extrapolation) voor, een framework dat LLM's trainen om kennis te extrapoleren over gestructureerde kennisgrafieken (KG's) via tool-augmented reinforcement learning (RL). In plaats van handmatig ontworpen prompts, leert het model autonoom wanneer en hoe het kennis moet ophalen en redeneren.

Kerncomponenten:

1. Agente Pipeline met drie beurten (Turns):
   Het model genereert een traject van drie beurten, afgewisseld met twee deterministische tool-aanroepen:

   • Beurt 1 (Entiteitsextractie): Het model analyseert de query en extraheren sleutelconcepten (bijv. ziekten, moleculen) binnen <extract_entities> tags.
   • Tool 1 (Entiteitsgroepconstructie): Een externe tool koppelt deze concepten aan de KG en bouwt groepen van semantisch vergelijkbare entiteiten (buren in de grafiek).
   • Beurt 2 (Groepfiltering): Het model filtert deze groepen op relevantie en selecteert de meest geschikte groepen binnen <filtered_groups> tags. Dit is cruciaal om ruis te verwijderen.
   • Tool 2 (Cross-group Triplet Retrieval): De tool haalt alle KG-triplets op die de geselecteerde groepen met elkaar verbinden.
   • Beurt 3 (Associatief Redeneren): Het model voert associatief redeneren uit. Als de opgehaalde triplets het antwoord niet direct bevatten, construeert het nieuwe triplets door semantisch vergelijkbare concepten uit de groepen te substitueren (bijv. het afleiden van (melatonine, behandelt, slapeloosheid) uit het patroon (hormoon, behandelt, mentale stoornis)).

2. Tool-Interface:
   De tools zijn deterministisch en extern aan het beleid (policy). Ze worden getriggerd door speciale tokens die het model genereert.

   • Tool 1: Linkt query-entiteiten aan KG-nodes en bouwt semantische groepen.
   • Tool 2: Haalt triplets op die entiteiten uit verschillende groepen verbinden.

3. Training met Reinforcement Learning (RL):

   • Algoritme: Het framework gebruikt GRPO (Group Relative Policy Optimization) binnen het VERL-framework.
   • End-to-End Optimalisatie: De hele multi-turn traject (extractie, filtering, redeneren) wordt gezamenlijk geoptimaliseerd. Tool-outputs worden niet getraind (ze zijn deterministisch), maar de beloningssignalen stromen terug naar de modelbeslissingen.
   • Curriculum Beloning: Een drie-fasen beloningsfunctie wordt gebruikt:
     1. Fase 1: Alleen formatbeloning (leert het model de structuur van de output).
     2. Fase 2: Combinatie van format en nauwkeurigheid.
     3. Fase 3: Alleen nauwkeurigheid (format is dan geautomatiseerd).
   • Er is geen Supervised Fine-Tuning (SFT) vooraf; het model leert het gedrag puur via RL-exploratie.

Belangrijkste Bijdragen

1. Leren van Kennisextrapolatie: SAKE demonstreert dat het vermogen om kennis te extrapoleren via analogie en associatief redeneren een leerbare vaardigheid is die kleine, open-weight modellen kunnen verwerven zonder complexe prompting.
2. Unificatie van Ophalen en Redeneren: In tegenstelling tot eerdere werken die ophalen en redeneren scheiden of handmatig ontwerpen, unificeert SAKE dit in één end-to-end beleid dat via RL wordt geoptimaliseerd.
3. Efficiëntie: Het framework vereist slechts één inferentiepass (single inference pass), in tegenstelling tot de iteratieve, multi-call benaderingen van bestaande agentic frameworks.
4. Robuustheid over Schalen: Het werkt effectief op zowel kleine (biomedische) als zeer grote (commonsense) kennisgrafieken.

Resultaten

De experimenten zijn uitgevoerd op twee benchmarks met verschillende schalen:

• Biomedisch (UMLS KG, 135 nodes): PubMedQA, BioASQ, ProcessBank.
• Commonsense (ConceptNet, 844k nodes): CommonsenseQA.

Kernbevindingen:

• Prestaties: Een met SAKE gefinetuned Qwen2.5-7B model presteert beter dan GPT-3.5-Turbo (gecombineerd met state-of-the-art frameworks zoals GIVE) op zowel biomedische (75,4% vs 70,1%) als commonsense benchmarks (81,3% vs 74,7%).
• Schaalbaarheid: SAKE verbetert de prestaties consistent op zowel 3B als 7B modellen, terwijl andere methoden (zoals SFT of GRPO zonder tools) onstabiele of lagere winsten laten zien.
• Efficiëntie: SAKE vermindert het token-gebruik met meer dan 90% vergeleken met frameworks zoals ToG en GIVE, omdat het geen iteratieve LLM-calls vereist.
• Ablatie-studies:
  • Het verwijderen van de groepfiltering leidt tot een daling in nauwkeurigheid (ruis in de context).
  • Agente ophaling (leren wat te ophalen) presteert significant beter dan vooraf berekende (pre-computed) ophaling.
  • Kennisextrapolatie (het construeren van nieuwe triplets) is cruciaal, vooral in onvolledige KG's (zoals BioASQ), waar directe antwoorden vaak ontbreken.

Betekenis en Impact

SAKE bewijst dat complexe redeneertaken in kennisintensieve domeinen niet afhankelijk hoeven te zijn van enorme modellen of ingewikkelde prompting-engineering. Door het combineren van gestructureerde kennisgrafieken met reinforcement learning, kunnen kleinere, open-weight modellen leren om effectief te extrapoleren en gaten in externe kennis te overbruggen. Dit opent de deur voor het inzetten van kostenefficiënte, betrouwbare AI-systemen in gespecialiseerde domeinen (zoals de gezondheidszorg) waar het onderhouden van complete kennisbanken onpraktisch is. De code is open-source beschikbaar gesteld.