Sintesi Tecnica: Assign and Add: Uno Studio Meccanicistico dell'Aritmetica Composizionale

Problema

I grandi modelli linguistici (LLM) dimostrano la capacità di comporre abilità per eseguire compiti complessi mai visti durante l'addestramento, eppure i meccanismi specifici che abilitano questa "generalizzazione composizionale" rimangono elusivi. Sebbene la generalizzazione sia spesso attribuita alla composizione di meccanismi atomici, l'interazione tra meccanismi distinti — specificamente il binding delle variabili (richiamo associativo tramite induction heads) e l'aritmetica modulare (rappresentazioni strutturate tramite feature di Fourier) — non è stata studiata sistematicamente in un contesto controllato.

Questo articolo affronta la domanda: Come imparano i transformer a comporre il binding delle variabili e i meccanismi di addizione modulare per generalizzare oltre la distribuzione di addestramento, e come queste capacità emergono durante l'addestramento?

Metodologia

Gli autori impiegano un setup sperimentale controllato che coinvolge un transformer piccolo e meccanicisticamente interpretabile, addestrato su un compito sintetico che combina l'assegnazione di variabili e l'addizione modulare.

• Architettura del Modello: Un transformer a 2 strati e singolo head con $d_{model} = d_{head} = 128$ e $d_{mlp} = 512$. L'architettura omette deliberatamente bias, strati di normalizzazione e l'MLP dello strato 1 per semplificare la scoperta dei circuiti. Utilizza embedding posizionali assoluti in stile GPT-2 e masking causale.
• Definizione del Compito: Il modello è addestrato per eseguire l'addizione modulare ($N=59$) dove gli operandi possono essere costanti o variabili assegnate dinamicamente nel contesto (ad es. c=1, b=17, a=42, b+19=?).
• Distribuzione dei Dati: Le sequenze sono tokenizzate con padding casuale per evitare l'affidamento alle posizioni assolute. Il set di addestramento è partizionato per valutare la generalizzazione:
  • 0-variabili (0-var): Pura addizione di costanti.
  • 1-variabile (1-var) e 2-variabili (2-var): Addizione che coinvolge variabili assegnate.
  • Set di Esclusione (Held-out Sets): Circa il 30% delle coppie di addizione viene rimosso dall'addestramento (ad es. se 1+3 non è visto, tutte le sequenze che lo contengono vengono escluse). Inoltre, specifiche variabili sono limitate in certe posizioni nel set di addestramento per testare la generalizzazione posizionale.
• Tecniche di Analisi: Lo studio utilizza strumenti di interpretabilità meccanicistica (tramite la libreria TransformerLens) per analizzare i pattern di attenzione, i vettori dello stream residuo e le matrici dei pesi. Impiega inoltre probe lineari e analisi del gradiente teorico per tracciare la dinamica dell'addestramento.

Contributi Chiave

1. Modello Trattabile per l'Aritmetica Composizionale: Gli autori presentano un modello semplificato che permette l'interpretazione meccanicistica di come l'assegnazione delle variabili si componga con l'aritmetica modulare.
2. Dinamiche di Addestramento Disaccoppiate: L'analisi rivela che l'aritmetica modulare viene appresa prima e indipendentemente dal binding delle variabili. Il modello raggiunge un'accuratezza quasi perfetta sui compiti 0-var pur imparando ancora a eseguire il binding delle variabili.
3. Spiegazione Meccanicistica della Generalizzazione: Il documento identifica che la generalizzazione deriva dalla composizione di due circuiti specifici:
   • Una induction head modificata nello Strato 1 che recupera le variabili assegnate e instrada i loro valori.
   • Un circuito di addizione modulare nello Strato 2 (specificamente la matrice Output-Value e l'MLP dello Strato 2) che esegue l'aritmetica sui valori recuperati.
4. Framework Teorico: Gli autori forniscono un'analisi teorica dimostrando come la dinamica del gradiente faciliti la composizione dei circuiti di binding delle variabili e di addizione modulare, mostrando che il gradiente negativo sulla matrice Query-Key dello Strato 2 porta naturalmente alla formazione di un blocco di memoria associativa.

Risultati

Dinamiche di Addestramento e Fasi

Il processo di addestramento presenta tre fasi distinte:

1. Apprendimento dell'Addizione Modulare: Il modello impara prima a eseguire l'addizione su sequenze senza variabili (0-var), raggiungendo un'alta accuratezza prima che il binding delle variabili sia pienamente stabilito.
2. Apprendimento dell'Assegnazione delle Variabili: Il modello apprende successivamente la struttura necessaria per l'assegnazione delle variabili. Uno "spike" netto nell'accuratezza (intorno al passo 5.300–5.400) coincide con l'emergere del circuito di routing dell'assegnazione delle variabili, dove l'attenzione dello Strato 1 instrada correttamente le costanti alle loro variabili assegnate.
3. Fase di Raffinamento: Una seconda fase (passi 14.500–16.000) prevede il raffinamento del routing per sequenze "difficili" (ad es. variabili specifiche in posizioni escluse), correggendo i fallimenti iniziali nel prestare attenzione a determinate posizioni delle variabili.

Capacità di Generalizzazione

• Generalizzazione Posizionale: Il modello generalizza a nuovi ordini di variabili (ad es. predire b+14 dopo aver visto 14+b).
• Generalizzazione Out-of-Distribution (OOD): Il modello predice con successo coppie di addizione escluse (ad es. 1+3) che non sono mai state viste in alcuna forma durante l'addestramento, raggiungendo un'accuratezza del 99,4% su set di test curati.
• Requisiti di Dati: La generalizzazione a coppie di costanti non viste richiede l'addestramento su almeno il ~25% della distribuzione totale ($f \approx 0.25$). La generalizzazione a sequenze a 2 variabili richiede una frequenza relativa di circa il ~20% ($r \approx 0.2$) di sequenze a 2 variabili nei dati di addestramento.

Risultati Meccanicistici

• Rappresentazione Condivisa: I vettori dello stream residuo per l'addizione con variabile assegnata e l'addizione di pura costante sono quasi identici (similarità del coseno $\approx 1$) prima dell'MLP dello Strato 2, indicando un meccanismo computazionale condiviso.
• Feature di Fourier: Le preattivazioni dei neuroni dell'MLP dello Strato 2 corrispondono alla forma teorica delle feature di Fourier ($\cos(\frac{2\pi k}{N}n + \alpha_k) + \cos(\frac{2\pi k}{N}m + \beta_k)$), confermando l'uso di meccanismi periodici stabiliti per l'aritmetica modulare.
• Modifica dell'Induction Head: L'attenzione dello Strato 1 agisce come un "previous-token head" per le costanti, prestando attenzione alla variabile immediatamente precedente. La matrice Query-Key dello Strato 2 forma quindi un forte pattern di identità sugli embedding delle variabili, permettendo al token = di prestare attenzione direttamente alle costanti corrette associate alle variabili.

Significato

L'articolo sostiene che la generalizzazione composizionale nei transformer può essere una naturale conseguenza della compositività dei meccanismi interni. Dimostrando che meccanismi distinti (induction heads e aritmetica basata su Fourier) possono essere appresi indipendentemente e poi composti tramite routing dell'attenzione, questo lavoro fornisce una base meccanicistica per comprendere come i transformer generalizzino a compiti di ragionamento complessi.

Gli autori mantengono la modestia riguardo al proprio ambito, notando che lo studio è limitato a un contesto piccolo e controllato per facilitare l'analisi. Non pretendono che questi risultati spieghino direttamente tutti i comportamenti degli LLM, ma suggeriscono che offrano un framework per comprendere come le capacità di generalizzazione emergano nei transformer, potenzialmente motivando future investigazioni su compiti più complessi e leggi di scala.