Beyond Black-Box Interventions: Latent Probing for Faithful Retrieval-Augmented Generation

Probleemstelling: Contextuele Trouw in RAG

Retrieval-Augmented Generation (RAG) systemen zijn ontworpen om Large Language Models (LLM's) te verrijken met externe, actuele kennis. Een kritieke uitdaging in de praktijk is echter het behoud van contextuele trouw (contextual faithfulness). Bestaande systemen falen vaak wanneer ze geconfronteerd worden met:

1. Kennisconflicten: Wanneer de geretrieveerde context in strijd is met de interne parametrische kennis van het model.
2. Contextuele ruis: Wanneer de context irrelevante informatie bevat die de aandacht van het model afleidt.

Bestaande methoden om dit op te lossen (zoals speciale prompts, decoding-calibratie of voorkeursoptimalisatie zoals DPO) behandelen het LLM als een zwarte doos. Ze interveniëren extern zonder inzicht te hebben in het interne redeneerproces. Dit leidt tot methoden die kwetsbaar zijn, veel data vereisen en niet kunnen diagnosticeren waarom een model faalt bij specifieke conflicten.

Methodologie: ProbeRAG Framework

De auteurs introduceren ProbeRAG, een raamwerk dat de "zwarte doos" opent door de interne latent space van het model te analyseren en te sturen. Het raamwerk bestaat uit drie fasen:

1. Fijnkorrelige Kennis-Pruning (Fine-Grained Knowledge Pruning)

Om de invloed van ruis te minimaliseren, wordt de geretrieveerde context niet als één blok behandeld, maar gedecomposeerd in onafhankelijke, atomische zinnen (kennisstatements).

• Een externe LLM (bijv. GPT-4o) breekt de context op in feitelijke statements.
• Vervolgens wordt de semantische gelijkenis tussen deze statements en de query berekend.
• Alleen de top-k meest relevante statements worden behouden, waardoor irrelevante ruis wordt gefilterd voordat de conflictanalyse plaatsvindt.

2. Latent Conflict Probing

Dit is de kerninnovatie van het artikel. De auteurs hebben ontdekt dat kennis die in overeenstemming is met het model (aligned) en kennis die in conflict is (conflicting), lineair scheidbaar zijn in de latent space (verborgen staten) van het model.

• Observatie: Contextuele ruis verhoogt de entropie van deze representaties, waardoor het onderscheid vervaagt.
• Implementatie: Een lichtgewicht "probe" (een MLP-classificator) wordt getraind op de verborgen staten van het model om te voorspellen of een kennisstatement in conflict is met de interne kennis van het model.
• Markering: Als een statement als conflict wordt gedetecteerd, wordt dit expliciet gemarkeerd in de input met speciale tokens (bijv. <conflict> en </conflict>).

3. Conflict-bewuste Aandacht (Conflict-Aware Attention)

Om ervoor te zorgen dat het model daadwerkelijk naar de gemarkeerde conflicten luistert in plaats van terug te vallen op zijn interne kennis, wordt een extra attention-guidance loss toegevoegd aan het trainingsdoel.

• De loss functie straft het model af als het weinig aandacht besteedt aan de tokens die als conflict zijn gemarkeerd.
• Het totale trainingsdoel is een gewogen som van de standaard taalmodeling loss en deze attention loss:
  $$L_{SFT} = (1 - \lambda)L_{CE} + \lambda L_{Attn}$$
• Dit dwingt het model om tijdens de inferentie meer computationele resources te alloceren aan de externe, conflicterende bewijslast.

Belangrijkste Bijdragen

1. Interne Analyse: De auteurs onthullen dat conflicterende en gealigneerde kennisstaten lineair scheidbaar zijn in de latent space van LLM's, en dat ruis systematisch de entropie van deze representaties verhoogt.
2. Nieuw Raamwerk (ProbeRAG): Een methode die verschuift van externe correcties naar interne begeleiding door gebruik te maken van latent probing en attention-modulatie.
3. Superieure Prestaties: Het raamwerk demonstreert consistent betere resultaten dan state-of-the-art baselines op meerdere benchmarks, zonder de interne kennis van het model te hoeven "wissen" (wat vaak leidt tot catastrofale vergetelheid).

Resultaten

Het team heeft ProbeRAG geëvalueerd op drie modellen (LLaMA-3.1-8B, Qwen3-8B, Mistral-7B) en meerdere datasets (FaithEval, ConFiQA, SQuAD).

• Prestaties: ProbeRAG behaalde State-of-the-Art (SOTA) resultaten. Bijvoorbeeld, op de FaithEval dataset verbeterde het de F1-score van de beste concurrent (CANOE) met ongeveer +3% en de Exact Match (EM) met +8% voor LLaMA-3.1-8B.
• Robuustheid: Het systeem presteerde sterk in scenario's met zowel feitelijke als logische kennisconflicten.
• Ablatiestudies:
  • Het verwijderen van de Conflict Probing leidde tot de grootste daling in prestaties, wat aantoont dat het detecteren van conflicten cruciaal is.
  • Het verwijderen van de Knowledge Pruning resulteerde in een afname van de nauwkeurigheid, wat de noodzaak van ruisfiltering bevestigt.
  • Er werd een positieve correlatie gevonden tussen de nauwkeurigheid van de conflict-detectie (de probe) en de uiteindelijke kwaliteit van het gegenereerde antwoord.
• Attention Visualisatie: Heatmaps tonen aan dat het getrainde model significant meer aandacht besteedt aan de tokens die als conflict zijn gemarkeerd, vergeleken met het oorspronkelijke model.

Significantie en Impact

Deze studie is significant omdat het een fundamentele verschuiving biedt in hoe we RAG-systemen verbeteren:

• Van Black-Box naar White-Box: In plaats van te gokken met prompts of decoding-strategieën, analyseert en sturen de auteurs het interne redeneerproces van het model.
• Behoud van Interne Kennis: In tegenstelling tot methoden die het model herschrijven (knowledge editing), "gids" ProbeRAG het model om de context te vertrouwen zonder de onderliggende parametrische kennis te verliezen.
• Generalisatie: De methode werkt effectief over verschillende modelarchitecturen en schalen, wat suggereert dat het fenomeen van lineair scheidbare conflicten een universeel kenmerk is van moderne LLM's.

Kortom, ProbeRAG biedt een principieel en effectief pad naar betrouwbaardere, minder hallucinerende RAG-systemen door inzicht te nemen in en te sturen op de interne dynamiek van kennisconflicten.