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

Sintesi Tecnica: Un circuito di rilegamento condizionato dal recupero per il tracciamento dinamico delle entità nei Large Language Models

Problema
I Large Language Models (LLM) devono interpretare correttamente il contesto e recuperare informazioni rilevanti legando le entità ai loro attributi e aggiornando questi legami al variare degli stati. Sebbene studi precedenti abbiano dimostrato che gli LLM possono eseguire il tracciamento delle entità, il loro comportamento rimane inconsistente tra diversi contesti, e i meccanismi computazionali sottostanti per gestire le transizioni di stato esplicite in contesti naturalistici non sono ancora pienamente compresi. Nello specifico, non è chiaro se i modelli mantengano un "modello di situazione" aggiornato globalmente dopo ogni cambiamento di stato (ad esempio, un'operazione di scambio o swap) o se impieghino un meccanismo più mirato per aggiornare i legami solo quando il recupero (retrieval) è richiesto.

Metodologia
Gli autori investigano questo problema utilizzando un compito di tracciamento dello stato dinamico che coinvolge scatole e oggetti, adattato da Kim e Schuster (2023). Lo studio impiega interventi causali sulle attivazioni di quattro modelli istruiti tramite instruction-tuning: Gemma2-9B-it, Gemma3-12B-it, Llama3.2-3B-it e Llama3.1-8B-it. L'analisi si articola in quattro fasi:

1. Analisi di mediazione causale: Utilizzando interventi di scambio (interchange interventions), gli autori tracciano il flusso di informazioni a livello di token. Costruiscono input controfattuali per determinare come l'informazione associata alle scatole di riferimento, alle scatole target dello scambio e agli oggetti target influenzi la previsione finale.
2. Path Patching: Per localizzare il meccanismo a livello di testa di attenzione (attention head), gli autori applicano il path patching (V-patching e Q-patching) per identificare un circuito compatto. Selezionano iterativamente i percorsi che impattano più negativamente sulle prestazioni quando vengono corrotti, isolando un sottoinsieme di teste di attenzione responsabili del compito.
3. Interventi di scambio specifici per il ruolo: Gli autori testano i ruoli funzionali dei gruppi di teste identificati effettuando il patching dei pattern di attenzione tra le esecuzioni originali e quelle controfattuali. Questo distingue tra le teste che mediano la selezione dell'indirizzo (routing) e quelle che trasferiscono il contenuto dell'oggetto.
4. Interventi sul Binding ID: Adattando il framework di Feng e Steinhardt (2024), gli autori eseguono interventi causali nei sottospazi Query (Q) e Key (K) di specifici gruppi di teste. Stimano le direzioni del binding ID e i vettori di spostamento per testare se il routing del recupero si basi sul matching Q/K degli identificatori di legame (binding identifiers).

Contributi Chiave e Risultati

• Evidenza del rilegamento condizionato dal recupero (Retrieval Conditioned Rebinding): L'analisi di mediazione causale rivela che l'informazione relativa alle scatole di riferimento e alle scatole target dello scambio non fluisce direttamente dalle loro menzioni iniziali, ma emerge dopo l'operazione di scambio e si propaga verso la previsione. Al contrario, l'informazione sull'oggetto viene instradata in modo più diretto. Questa dissociazione supporta l'ipotesi del "Rilegamento Condizionato dal Recupero": il modello non ricodifica globalmente l'intero stato post-scambio. Invece, preserva i legami originali e applica una rimappatura sui binding ID solo durante il recupero della risposta.
• Identificazione di un circuito compatto: Il path patching identifica un circuito di attenzione specifico composto dal 3-10% dei componenti completi del modello (ad esempio, 38 teste in Gemma-9B) che recupera prestazioni quasi pari al modello completo (0.89 di accuratezza rispetto a 1.00 del modello completo). Questo circuito consiste in cinque gruppi funzionali:
  • Gruppo E (Ancora del Legame - Binding Anchor): Scrive l'informazione di identità della scatola.
  • Gruppo D (Trasmettitore della Posizione di Scambio - Swap Position Transmitter): Propaga l'informazione dello scambio.
  • Gruppo C (Aggiornatore di Posizione - Position Updater): Modella l'attenzione verso le posizioni di legame aggiornate.
  • Gruppo B (Dereferenziatore - Dereferencer): Agisce come un puntatore, mediando la selezione dell'indirizzo.
  • Gruppo A (Recuperatore della Risposta - Answer Retriever): Recupera il token dell'oggetto target.
• Meccanismo di Routing dell'Indirizzo: Gli interventi di scambio confermano che il Gruppo B funziona come un puntatore, sensibile alla selezione dell'indirizzo. Il patching dei suoi pattern di attenzione sposta il focus del modello verso la corretta posizione sorgente senza trasferire il contenuto dell'oggetto controfattuale.
• Firme nei sottospazi Q/K: Gli interventi sul binding ID rivelano che il meccanismo di puntamento si basa sul matching Q/K, ma l'implementazione varia a seconda della famiglia di modelli:
  • Modelli Gemma: Mostrano una chiara firma di matching Q/K in cui l'informazione di legamento è espressa in entrambi i sottospazi query e key. Spostare il vettore di query reindirizza efficacemente l'attenzione.
  • Modelli Llama: Si affidano principalmente ai vettori key per l'informazione di legame. Gli interventi sul lato key producono gli spostamenti di attenzione più grandi, mentre gli effetti sul lato query sono più deboli.
• Generalizzabilità: Il meccanismo di rilegamento condizionato dal recupero è osservato nelle famiglie Gemma e Llama, suggerendo che sia una strategia computazionale condivisa per il tracciamento dinamico dello stato, nonostante le differenze nelle specifiche firme del circuito.

Significatività e Rivendicazioni
Il documento sostiene di fornire una spiegazione interpretabile, a livello di circuito, di come gli LLM eseguano il tracciamento dello stato dipendente dal contesto. I risultati suggeriscono che gli LLM non mantengono necessariamente una simulazione mentale completa e globalmente aggiornata del mondo dopo ogni cambiamento di stato. Invece, utilizzano una strategia di "rilegamento" in cui l'informazione relativa allo scambio viene preservata e applicata localmente al momento del recupero per reindirizzare l'attenzione verso l'oggetto corretto tramite i binding ID aggiornati.

Gli autori osservano che, sebbene la strategia computazionale (rilegamento condizionato dal recupero) appaia coerente tra le famiglie di modelli, l'implementazione a livello di circuito (specificamente la dipendenza dai sottospazi Q vs K) differisce. Ciò implica che il meccanismo può essere una soluzione condivisa al problema del tracciamento dinamico, ma la sua realizzazione dipende dalle specifiche architetture dei modelli o dalle distribuzioni di pre-training.

Limitazioni
Gli autori riconoscono che il loro setting sperimentale è altamente controllato (singolo scambio, oggetti a singolo token, scatole nominate alfabeticamente) e potrebbe non catturare l'intera gamma di meccanismi nel discorso naturale. Inoltre, le direzioni del binding ID utilizzate negli interventi sono stimate sotto un'approssimazione di rappresentazione lineare, il che potrebbe sottostimare i meccanismi se l'informazione di legamento è rappresentata in modo non lineare o distribuita su più sottospazi.