Sintesi Tecnica: Pruning Multi-Granulare dei Nodi per la Scoperta di Circuiti Causali

1. Definizione del Problema

La scoperta di circuiti causali mira a identificare sottoreti minimali all'interno dei Large Language Models (LLM) che guidano causalmente comportamenti specifici. Sebbene i metodi esistenti abbiano fatto progressi nella interpretabilità meccanicistica, essi affrontano due limitazioni primarie:

1. Scalabilità Computazionale: Gli approcci attuali si basano prevalentemente sul pruning degli archi (rimozione delle connessioni tra componenti). Poiché il numero di archi in un transformer cresce quadraticamente con il numero di nodi ($O(|N|^2)$), il pruning a livello di arco diventa computazionalmente intrattabile all'aumentare delle dimensioni del modello. Inoltre, questi metodi richiedono spesso di memorizzare rappresentazioni dello stream residuo e attivazioni intermedie disgiunte, portando a un uso della memoria proibitivo.
2. Limitazioni di Granularità: I metodi esistenti operano tipicamente a una granularità grossolana, trattando interi head di attenzione o blocchi MLP come unità atomiche. Ciò trascura la possibilità che specifici neuroni (nodi) all'interno di questi componenti siano responsabili di comportamenti target, risultando in circuiti recuperati che mancano di strutture fini e mantengono componenti non necessari.

2. Metodologia

Gli autori propongono un framework di pruning dei nodi multi-granulare che scopre circuiti minimali attraverso molteplici livelli di granularità all'interno di un unico step di ottimizzazione unificato.

2.1 Design della Granularità Gerarchica

Invece di eseguire il pruning degli archi tra componenti grossolani, il framework apprende maschere su nodi a quattro distinti livelli:

• Livello di Blocco (Block Level): Interi Blocchi di Attenzione e Blocchi MLP.
• Livello di Head (Head Level): Singoli Attention Head.
• Livello di Neurone (Neuron Level): Singoli neuroni all'interno dei layer di attenzione e dei due layer lineari dei blocchi MLP (MLP1/Keys e MLP2/Values).

2.2 Setup di Pruning Causale a Due Stream

Per identificare l'importanza causale senza introdurre shift di distribuzione (un rischio con il semplice zero-ablation), il metodo impiega un passaggio forward a due stream:

• Stream Pulito (Clean Stream): Il modello elabora un input progettato per elicitare un comportamento specifico (es. predire un anno > 50).
• Stream Corrotto (Corrupted Stream): Il modello elabora una versione minimamente perturbata dell'input (es. invertire i pronomi di genere o cambiare l'anno di partenza) per elicitare una predizione differente.
• Interpolazione: Per ogni nodo $n_i$, una maschera apprendibile $m_i \in [0, 1]$ interpola tra l'attivazione pulita ($h^{clean}$) e l'attivazione corrotta ($h^{corrupted}$):
  $$h^{node}_i = m_i \cdot h^{clean}_i + (1 - m_i) \cdot h^{corrupted}_i$$
  L'ottimizzazione cerca di trovare maschere che permettano al modello di produrre l'output corretto (corrispondente al comportamento pulito) massimizzando al contempo la sparsità.

2.3 Ottimizzazione e Parametrizzazione delle Maschere

• Distribuzione Hard Concrete: Le maschere sono parametrizzate utilizzando la distribuzione Hard Concrete per consentire l'ottimizzazione basata sul gradiente pur inducendo un comportamento approssimativamente binario (0 o 1).
• Obiettivo Unificato: L'obiettivo di training combina tre componenti:
  1. Loss di Fedeltà ($L_{faith}$): Una divergenza KL che assicura che la distribuzione dell'output del circuito sottoposto a pruning corrisponda a quella del modello base.
  2. Loss del Task ($L_{task}$): Una loss di ranking basata sul margine che assicura che il circuito distingua correttamente tra le opzioni reali e quelle distratrici.
  3. Regolarizzatore di Sparsità ($L_0$): Incoraggia le maschere a essere zero, minimizzando la dimensione del circuito.
• Consistenza Gerarchica: Post-training, le maschere vengono binarizzate con un vincolo che assicura che se un nodo genitore (es. un blocco MLP) viene rimosso, tutti i suoi nodi figli (es. i neuroni all'interno di quel blocco) siano deattivati.

3. Contributi Chiave

1. Framework Multi-Granulare: L'introduzione di un framework di ottimizzazione unificato che apprende maschere simultaneamente attraverso i blocchi del transformer, gli attention head, i blocchi MLP e i singoli neuroni.
2. Scalabilità ed Efficienza: Il metodo evita la scalabilità quadratica del pruning degli archi e la necessità di memorizzare le attivazioni intermedie per tutti gli archi. Riduce l'uso della memoria GPU di 3× fino a 11× rispetto ai metodi precedenti, consentendo la scoperta di circuiti su Llama 3.1-8B utilizzando una singola GPU da 40GB.
3. Scoperta Fine-Grained: L'approccio identifica circuiti significativamente più compatti rispetto ai baseline, raggiungendo specificamente un'alta sparsità ai livelli di blocco MLP e neurone, dove i metodi precedenti fallivano ampiamente.

4. Risultati Sperimentali

Il metodo è stato valutato su tre task standard di scoperta di circuiti (Indirect Object Identification, Gendered Pronouns e Greater Than) utilizzando modelli GPT-2, GPT-2-XL e Llama.

4.1 Sparsità e Compattezza

• Pruning di MLP e Neuroni: Il metodo proposto supera significativamente i baseline di edge-pruning (EAP, EP) e il pruning non strutturato dei pesi (DiscoGP). Nella configurazione meno favorevole, ha rimosso il 33,34% in più di blocchi MLP e il 59,8% in più di neuroni rispetto ai metodi precedenti.
• Strutture Dipendenti dal Task: I circuiti scoperti esibiscono topologie distinte in base al task:
  • IOI: Dominato da computazione MLP distribuita; l'attenzione è ampiamente rimossa tranne nei layer finali.
  • Gendered Pronouns (GP): Altamente localizzato e discontinuo; la computazione MLP è ristretta a pochi layer, mentre l'attenzione è sparsa.
  • Greater Than (GT): Esibisce una sparsità estrema con blocchi contigui di layer inattivi e pattern di esecuzione di tipo "skip-like".

4.2 Fedeltà e Consistenza

• Performance: I circuiti sottoposti a pruning mantengono un'alta accuratezza del task (spesso >95%) e una bassa divergenza KL rispetto al modello completo.
• Consistenza: Gli esperimenti su cinque seed casuali mostrano che i circuiti scoperti (metriche di fedeltà e profili di sparsità) sono altamente stabili.
• Analisi Qualitativa: In Llama 3.2-1B, gli attention head sopravvissuti nel task IOI sono stati trovati capaci di svolgere ruoli funzionali specifici (IO-Movers, S-Inhibitors), rispecchiando i circuiti scoperti manualmente in GPT-2, validando la capacità del metodo di recuperare insight meccanicistici.

4.3 Efficienza Computazionale

• Memoria: Il metodo richiede almeno 3× meno memoria rispetto ai baseline di edge-pruning (es. 6,27 GB vs 72,79 GB per EAP su GPT-2).
• Tempo: Il training si completa in meno di 6 minuti, significativamente più veloce dei metodi di edge-pruning (circa 50 minuti), sebbene leggermente più lento del metodo EAP a singolo passaggio (che scambia il tempo per un enorme uso di memoria).

5. Significato e Rivendicazioni

Il paper sostiene che la sua importanza primaria risieda nel superamento dei colli di bottiglia di scalabilità e granularità degli attuali metodi di scoperta dei circuiti. Passando dal pruning degli archi al pruning multi-granulare dei nodi, gli autori dimostrano che:

• Molti neuroni mantenuti dai metodi a granularità grossolana sono superflui e possono essere rimossi con un impatto trascurabile sulle performance.
• Modelli più grandi (es. Llama 8B) supportano circuiti altrettanto forti o più forti rispetto ai modelli più piccoli, suggerendo che l'aumento di scala introduce ridondanza piuttosto che richiedere circuiti task-specific più densi.
• Il metodo fornisce una via computazionalmente fattibile per scoprire circuiti causali fine-grained in modelli su larga scala, rivelando che comportamenti differenti sono codificati da pattern distinti di attività di attenzione, MLP e neuroni.

Gli autori riconoscono un limite: il framework identifica quali nodi sono necessari ma non recupera esplicitamente la struttura di interazione (archi) tra i nodi sopravvissuti. Suggeriscono come direzione futura la combinazione della selezione dei nodi con il recupero degli archi per inferire le strutture di interazione.