A retrieval conditioned rebinding circuit for dynamic entity tracking in large language models

Technische Samenvatting: Een Retrieval Conditioned Rebinding Circuit voor Dynamische Entiteits-tracking in Large Language Models

Probleemstelling
Large Language Models (LLM's) moeten context correct interpreteren en relevante informatie ophalen door entiteiten te koppelen aan hun attributen en deze koppelingen bij te werken wanneer staten veranderen. Hoewel eerder werk heeft aangetoond dat LLM's entiteit-tracking kunnen uitvoeren, blijft hun gedrag inconsistent over verschillende settings heen, en de onderliggende computationele mechanismen voor het afhandelen van expliciete staatsovergangen in natuurlijke contexten zijn niet volledig begrepen. Specifiek is het onduidelijk of modellen na elke staatsovergang (bijv. een swap-operatie) een globaal bijgewerkte "situatiemodel" bijhouden, of dat ze een meer gerichte methode hanteren om koppelingen alleen bij te werken wanneer retrieval vereist is.

Methodologie
De auteurs onderzoeken dit probleem met behulp van een taak voor dynamische staatstracking met dozen en objecten, aangepast van Kim en Schuster (2023). De studie maakt gebruik van causale interventies op de activaties van vier instructie-afgestemde modellen: Gemma2-9B-it, Gemma3-12B-it, Llama3.2-3B-it, en Llama3.1-8B-it. De analyse verloopt in vier stadia:

1. Causale Mediatie-analyse: Met behulp van interchange-interventies traceren de auteurs de informatiestroom op token-niveau. Ze construeren contrafeitelijke inputs om te bepalen hoe informatie geassocieerd met referentiedozen, swap-doeldozen en doelobjecten de uiteindelijke voorspelling beïnvloedt.
2. Path Patching: Om het mechanisme op het niveau van de attention head te lokaliseren, passen de auteurs path patching toe (V-patching en Q-patching) om een compact circuit te identificeren. Ze selecteren iteratief paden die de prestaties het meest negatief beïnvloeden wanneer ze worden gecorrumpeerd, waardoor een subset van attention heads wordt geïsoleerd die verantwoordelijk is voor de taak.
3. Rol-specifieke Interchange Interventies: De auteurs testen de functionele rollen van geïdentificeerde head-groepen door attention-patronen te patchen tussen originele en contrafeitelijke runs. Dit onderscheidt heads die adressselectie mediëren (routing) van heads die objectinhoud overdragen.
4. Binding ID Interventies: Geïnspireerd door het framework van Feng en Steinhardt (2024), voeren de auteurs causale interventies uit in de Query (Q) en Key (K) subspaces van specifieke head-groepen. Ze schatten binding ID-richtingen en shift-vectoren in om te testen of retrieval-routing vertrouwt op Q/K-matching van binding-identificatoren.

Belangrijkste Bijdragen en Resultaten

• Bewijs voor Retrieval Conditioned Rebinding: Causale mediatie-analyse laat zien dat informatie over de referentie- en swap-doeldozen niet direct voortkomt uit hun initiële vermeldingen, maar ontstaat na de swap-operatie en naar de voorspelling propageert. Daarentegen wordt objectinformatie directer gerouteerd. Deze dissociatie ondersteunt de "Retrieval Conditioned Rebinding"-hypothese: het model codeert niet globaal de volledige post-swap staat opnieuw. In plaats daarvan behoudt het oorspronkelijke bindings en past het een hermapping toe over binding ID's, enkel tijdens de antwoord-retrieval.
• Identificatie van een Compact Circuit: Path patching identificeert een specifiek attention head circuit bestaande uit 3–10% van de volledige modelcomponenten (bijv. 38 heads in Gemma-9B) dat bijna de volledige modelprestaties herstelt (0,89 nauwkeurigheid versus 1,00 voor het volledige model). Dit circuit bestaat uit vijf functionele groepen:
  • Groep E (Binding Anchor): Schrijft box-identiteitsinformatie.
  • Groep D (Swap Position Transmitter): Propageert swap-informatie.
  • Groep C (Position Updater): Vormt de aandacht richting bijgewerkte binding-posities.
  • Groep B (Dereferencer): Fungeert als een pointer, waarbij adressselectie wordt gemedieerd.
  • Groep A (Answer Retriever): Haalt de doel-object token op.
• Mechanisme van Adress-Routing: Interchange-interventies bevestigen dat Groep B fungeert als een pointer, gevoelig voor adressselectie. Het patchen van de attention-patronen van deze groep verschuift de focus van het model naar de juiste bronpositie zonder de contrafeitelijke objectinhoud over te dragen.
• Q/K Subspace Signatures: Binding ID-interventies onthullen dat het pointer-mechanisme vertrouwt op Q/K-matching, maar de implementatie varieert per modelfamilie:
  • Gemma Modellen: Vertonen een duidelijke Q/K-matching signature waarbij binding-informatie wordt uitgedrukt in zowel de query- als de key-subspaces. Het verschuiven van de query-vector herleidt de aandacht effectief.
  • Llama Modellen: Vertrouwen primair op key-vectoren voor binding-informatie. Key-zijde interventies produceren de grootste attention-verschuivingen, terwijl de effecten aan de query-zijde zwakker zijn.
• Generaliseerbaarheid: Het retrieval conditioned rebinding mechanisme wordt waargenomen in zowel de Gemma- als de Llama-families, wat suggereert dat het een gedeelde computationele strategie is voor dynamische staatstracking, ondanks verschillen in de specifieke circuit-signatures.

Betekenis en Claims
Het artikel claimt een interpreteerbare, circuit-niveau verklaring te bieden voor hoe LLM's context-afhankelijke staatstracking uitvoeren. De bevindingen suggereren dat LLM's niet noodzakelijkerwijs een volledige, globaal bijgewerkte mentale simulatie van de wereld bijhouden na elke staatswijziging. In plaats daarvan maken ze gebruik van een "rebinding"-strategie waarbij swap-gerelateerde informatie wordt bewaard en lokaal wordt toegepast tijdens de retrieval om de aandacht via bijgewerkte binding ID's naar het juiste object te leiden.

De auteurs merken op dat hoewel de computationele strategie (retrieval conditioned rebinding) consistent lijkt te zijn over modelfamilies, de implementatie op circuit-niveau (met name de afhankelijkheid van Q versus K subspaces) verschilt. Dit impliceert dat het mechanisme een gedeelde oplossing kan zijn voor het probleem van dynamische tracking, maar dat de realisatie ervan afhangt van specifieke modelarchitecturen of pretraining-distributies.

Beperkingen
De auteurs erkennen dat hun taaksetting zeer gecontroleerd is (enkele swap, single-token objecten, alfabetisch benoemde dozen) en mogelijk niet het volledige bereik van mechanismen in natuurlijke discours vangt. Daarnaast worden de binding ID-richtingen die in de interventies worden gebruikt, geschat onder een benaderde lineaire representatie-aanname, wat mechanismen kan onderschatten als binding-informatie niet-lineair wordt gerepresenteerd of verdeeld is over meerdere subspaces.