Dynamic Weight Grafting: Localizing Finetuned Factual Knowledge in Transformers

Il paper propone la "Dynamic Weight Grafting", una tecnica di analisi che grafta selettivamente pesi da un modello finetunato su uno pre-addestrato per dimostrare che le conoscenze fattoriali apprese vengono recuperate attraverso due percorsi distinti: un arricchimento del flusso residuo durante l'elaborazione dell'entità e un richiamo finale prima della generazione, localizzando questi meccanismi in specifiche componenti del modello come l'attenzione e le reti feedforward.

L'essenziale

🧠 Il Mistero: Dove vanno a finire i nuovi ricordi?

Immagina di avere un assistente personale molto intelligente (un Modello Linguistico o LLM) che ha letto milioni di libri e sa tutto sul mondo fino al 2023. Poi, gli dici: "Ehi, nel 2024 è stato eletto un nuovo Papa, Leo XIV".

Il modello impara questo fatto. Ma la domanda è: dove lo ha messo?

1. L'ha scritto subito nella "scheda" del Papa quando ha letto il nome?
2. Lo ha tenuto in un "foglio di appunti" nascosto e lo ha tirato fuori solo quando ha dovuto rispondere alla domanda?
3. O forse usa entrambi i metodi?

Fino a oggi, gli scienziati non potevano dirlo con certezza perché i loro strumenti per guardare dentro il cervello dell'AI erano troppo "grezzi": quando provavano a guardare un pezzo di memoria, spesso lo rompevano o lo cancellavano, rendendo impossibile capire come funzionava il resto.

🔧 La Nuova Soluzione: Il "Trapianto di Peso Dinamico"

Gli autori di questo studio hanno inventato un nuovo strumento chiamato Dynamic Weight Grafting (innesto dinamico dei pesi).

L'analogia perfetta:
Immagina che il modello AI sia un'orchestra.

• Il modello pre-addestrato (quello vecchio) è l'orchestra che suona musica classica.
• Il modello finetunato (quello che ha imparato il nuovo Papa) è la stessa orchestra che ha appena imparato una nuova canzone pop.

I vecchi metodi di ricerca erano come fermare l'orchestra a metà concerto, cambiare uno strumento e vedere cosa succede. Ma questo rovinava l'armonia precedente.

Il nuovo metodo, invece, è come un ingegnere del suono magico che, mentre l'orchestra suona, può:

• Sostituire solo il violino del 5° musicista con quello della versione "Pop".
• Sostituire solo il tamburo dell'ultimo musicista con quello della versione "Pop".
• E tutto il resto rimane l'orchestra "Classica".

In questo modo, possono vedere esattamente quale musicista (o quale parte del cervello) è responsabile per suonare la nota giusta del nuovo Papa.

🚀 Cosa hanno scoperto? Due strade per ricordare

Usando questo "ingegnere del suono", hanno scoperto che l'AI usa due strade diverse per ricordare un nuovo fatto, e a volte ne usa una sola, a volte entrambe:

1. La Strada dell'Arricchimento (Enrichment)

Quando l'AI incontra il nome "Papa Leo XIV" all'inizio della frase, la sua mente si "illumina" immediatamente. In quel preciso istante, il modello aggiunge tutte le informazioni sul nuovo Papa alla rappresentazione di quel nome.

• Metafora: È come se, appena vedi il nome "Papa" su un foglio, gli attaccassi subito un post-it con scritto "Leo XIV". Quando arriverà la fine della frase, il foglio avrà già l'informazione attaccata.

2. La Strada del Ricordo (Recall)

In altri casi, il modello non fa nulla quando vede il nome "Papa". Aspetta di arrivare all'ultima parola della frase (prima di rispondere). Solo in quel momento, come un detective che tira fuori la soluzione dal cappello, il modello "ricorda" e recupera l'informazione dal nulla.

• Metafora: È come se il foglio fosse vuoto mentre leggi, ma quando arrivi alla domanda "Chi è il Papa?", il modello fa un salto mentale e scrive "Leo XIV" solo in quel momento.

🔍 Il Dettaglio Tecnico (Ma semplificato)

Hanno anche guardato come avviene questo "Ricordo" finale. Hanno scoperto che non è magia, ma due ingranaggi specifici:

1. L'Attenzione (Attention): Funziona come gli occhi che guardano indietro. Cerca il nome del Papa all'inizio della frase.
2. La Rete di Alimentazione (Feedforward): Funziona come un archivio. Prende ciò che gli occhi hanno trovato e lo trasforma nella risposta corretta.

🎯 Perché è importante?

Prima, pensavamo che l'AI memorizzasse tutto in un unico posto o in modo confuso. Ora sappiamo che:

• L'AI è flessibile: a volte scrive subito il fatto, a volte lo aspetta alla fine.
• Possiamo "disattivare" una strada e l'altra funziona ancora (in parte).
• Se togliamo entrambe le strade, l'AI dimentica tutto e torna a essere come prima dell'addestramento.

In sintesi: Gli scienziati hanno costruito un "microscopio" che permette di vedere esattamente come un'intelligenza artificiale impara un nuovo fatto, senza distruggere la sua memoria. Hanno scoperto che l'AI ha due modi per imparare: o si prepara subito quando vede il nome, o si ricorda tutto all'ultimo secondo prima di parlare. E spesso, ha bisogno di entrambi i modi per essere perfetta!

Sommario tecnico

1. Il Problema

Quando un Large Language Model (LLM) apprende un nuovo fatto durante il fine-tuning (ad esempio, un nuovo papa eletto o un nuovo film), dove risiede fisicamente questa informazione all'interno dei parametri del modello?
Le domande chiave sono:

• Le entità vengono "arricchite" immediatamente con informazioni relazionali quando vengono elaborate?
• Le informazioni vengono recuperate "just-in-time" (appena prima della previsione) o sono memorizzate in modo ridondante?
• Quali componenti specifici del modello (attenzione, feed-forward network) sono responsabili del recupero di queste conoscenze?

Limiti degli approcci esistenti:
Le tecniche di localizzazione attuali, come il patching delle attivazioni (activation patching) o le ablazioni, presentano un difetto fondamentale: modificano o sostituiscono le attivazioni (lo stato residuo) in un punto specifico. Questo sovrascrive le informazioni calcolate a monte, rendendo impossibile distinguere se un componente sta estrando attivamente nuove informazioni o se sta semplicemente trasmettendo informazioni già elaborate in precedenza. Di conseguenza, non è possibile isolare con precisione i meccanismi necessari e sufficienti per il recupero della conoscenza.

2. Metodologia: Dynamic Weight Grafting

Gli autori propongono una nuova tecnica chiamata Dynamic Weight Grafting (innesto dinamico dei pesi) per colmare questo divario interpretativo.

• Concetto Base: Invece di manipolare le attivazioni, la tecnica sostituisce selettivamente sottoinsiemi di pesi (parametri) da un modello fine-tuned (SFT) in un modello pre-addestrato (Pretrained) durante la generazione del testo.
• Meccanismo: Il processo avviene token per token. Per ogni posizione del token $t$, è possibile scegliere dinamicamente quali componenti del modello (es. matrici di attenzione, reti feed-forward) utilizzare: quelli del modello pre-addestrato o quelli del modello fine-tuned.
• Vantaggio: Poiché i pesi vengono scambiati ma il flusso di calcolo precedente (lo stato residuo) rimane intatto (a differenza del patching delle attivazioni), questo metodo permette di testare quali componenti sono necessari e sufficienti per riprodurre il comportamento del modello fine-tuned senza distruggere le informazioni precedenti.

La metodologia distingue due tipi di innesto:

1. Position Grafting: Sostituisce tutti i pesi del modello fine-tuned per una specifica posizione del token (es. solo il primo token dell'entità o solo l'ultimo token).
2. Component Grafting: Sostituisce singoli componenti (es. solo la matrice di proiezione di output $O$ o solo i layer FFN) a specifiche posizioni.

3. Contributi Chiave e Risultati Sperimentali

Gli esperimenti sono stati condotti su quattro modelli (Llama3, Pythia, GPT-2 XL, Gemma) utilizzando dati sintetici (fatti su film e attori) e dati reali (articoli Wikipedia su film usciti dopo la data di release del modello).

A. Identificazione di Due Pathway di Recupero

L'analisi rivela che i modelli recuperano le conoscenze relazionali attraverso due percorsi distinti e spesso ridondanti:

1. Pathway di "Arricchimento" (Enrichment): Le informazioni relazionali vengono incorporate nello stato residuo quando viene elaborato il token dell'entità (es. "Zendaya"). I layer successivi possono estrarre queste informazioni da questo stato arricchito.
2. Pathway di "Recall" (Richiamo): Le informazioni vengono recuperate dinamicamente alla posizione dell'ultimo token prima della previsione, anche se l'entità non è stata "arricchita" in precedenza.

Risultato principale: Innestare i pesi del modello fine-tuned solo sulla prima entità (FE) oppure solo sull'ultimo token (LT) è sufficiente per recuperare quasi completamente le prestazioni del modello fine-tuned. Al contrario, se si innestano tutti i pesi tranne FE e LT, l'accuratezza crolla a livelli vicini allo zero (come il modello pre-addestrato).

B. Localizzazione dei Componenti nel Pathway di "Recall"

Utilizzando il Component Grafting, gli autori hanno mappato esattamente quali parti del modello sono responsabili del recupero delle informazioni:

• Attenzione Specifica del Task: È necessaria l'attenzione (meccanismi di attenzione) sui token della prima entità e sull'ultimo token. Questi meccanismi sono specifici del compito (la struttura della frase) e non necessariamente del fatto specifico.
• Estrazione Specifica della Relazione: Il recupero effettivo del fatto avviene nei layer finali (prima della previsione del token successivo) attraverso le matrici di proiezione di output ($O$) e le reti feed-forward (FFN).
  • Gli esperimenti mostrano che innestare solo la matrice $O$ e l'FFN dal modello fine-tuned (mantenendo l'attenzione del modello pre-addestrato o di un modello "task-specific") recupera gran parte delle prestazioni.
  • Questo suggerisce che l'attenzione prepara il contesto, mentre FFN e $O$ eseguono l'estrazione finale del fatto.

C. Generalizzazione

I risultati si mantengono coerenti sia su entità note che sconosciute e su dati strutturati (template) e non strutturati (articoli Wikipedia), sebbene il pathway di "recall" sia più forte nei modelli più recenti (Gemma, Llama3) rispetto a GPT-2 XL.

4. Significato e Implicazioni

• Superamento dei Limiti del Patching: Il Dynamic Weight Grafting offre un metodo di interpretazione meno distruttivo rispetto al patching delle attivazioni, permettendo di isolare i meccanismi causali senza sovrascrivere il flusso di informazioni a monte.
• Nuova Comprensione dell'Architettura: Dimostra che la conoscenza non è memorizzata in un unico luogo, ma è distribuita tra un'elaborazione iniziale (arricchimento dell'entità) e un recupero finale (recall). Inoltre, separa il ruolo dell'attenzione (che gestisce il contesto/task) da quello delle FFN (che estrae il fatto specifico).
• Implicazioni per il Knowledge Editing: La scoperta che i meccanismi di recupero possono essere localizzati in componenti specifici (come le FFN degli ultimi layer) potrebbe guidare tecniche più precise per l'editing della conoscenza nei LLM, permettendo di aggiornare fatti senza alterare la capacità generale del modello.
• Ridondanza: Il fatto che uno solo dei due pathway (arricchimento o recall) sia spesso sufficiente suggerisce che i modelli implementano strategie ridondanti per garantire l'affidabilità del recupero delle informazioni.

In sintesi, il paper fornisce una mappa meccanicistica dettagliata di come i Transformer integrano e recuperano nuove conoscenze fattuali, spostando il focus dall'analisi delle attivazioni all'analisi diretta dei parametri e dei loro ruoli funzionali specifici.