Understanding the Staged Dynamics of Transformers in Learning Latent Structure

1. Il Problema

Sebbene sia ben noto che i modelli Transformer possono scoprire strutture latenti dal contesto (ad esempio, gerarchie sintattiche o grafi causali), la dinamica di come acquisiscono queste diverse componenti rimane poco compresa. Esiste l'ipotesi che i modelli si limitino a "ricombinare" i dati di addestramento senza una vera generalizzazione. La difficoltà principale risiede nella complessità del linguaggio naturale, che rende difficile isolare e tracciare la progressione dell'apprendimento di strutture specifiche.

L'obiettivo di questo lavoro è investigare la dinamica di apprendimento delle strutture latenti in un ambiente controllato, analizzando se l'apprendimento avviene in modo continuo o attraverso fasi discrete (staged dynamics), e identificare i meccanismi sottostanti a queste transizioni.

2. Metodologia

Gli autori utilizzano il benchmark Alchemy (Wang et al., 2021), un compito di meta-apprendimento per rinforzo con una struttura latente ben definita e controllabile.

• Ambiente Alchemy: Il compito simula una "chimica" dove i giocatori applicano pozioni a una pietra per cambiarne le proprietà (colore, dimensione, rotondità, ricompensa). Le transizioni tra gli stati delle pietre formano un grafo diretto a struttura cubica con 8 vertici.

• Modello: Viene addestrato un piccolo Transformer decoder-only (circa 2 milioni di parametri) su tre varianti del compito, ciascuna progettata per testare abilità specifiche:

  1. Scoperta di pozioni nascoste (Withheld Potion Pair): Vengono rimossi dal set di supporto (support) tutti gli esempi relativi a una specifica coppia di pozioni. Il modello deve inferire l'effetto di queste pozioni mancanti basandosi sulle altre transizioni. Questo richiede di allineare due metà disconnesse del grafo chimico.
  2. Composizione (Composition): Il modello riceve tutte le transizioni a 1 passo (pozioni singole) e deve risolvere query multi-step (catene di 2-5 pozioni).
  3. Decomposizione (Decomposition): Il modello riceve sequenze complesse multi-step nel supporto e deve inferire l'effetto di una singola transizione (1-step) per risolvere una query.

• Analisi dei Fasi: Per comprendere la dinamica, gli autori scompongono l'accuratezza globale in componenti interpretabili basate su eventi probabilistici (es. $P[A]$: predizione corretta all'interno del supporto; $P[B|A]$: predizione della metà corretta del grafo; $P[C|A \cap B]$: corrispondenza esatta).

• Interventi Causali: Vengono utilizzati esperimenti di "congelamento" (freezing) dei parametri per identificare finestre di plasticità specifiche per ogni layer durante l'addestramento.

3. Contributi Chiave

Il lavoro fornisce una mappatura dettagliata di come un Transformer acquisisce abilità strutturate:

1. Apprendimento a Stadi Discreti: L'accuratezza non migliora linearmente, ma mostra plateau seguiti da salti improvvisi. Questi salti corrispondono all'acquisizione sequenziale di sottocompetenze interpretabili (es. prima restringere lo spazio delle ipotesi, poi identificare la metà corretta, infine la soluzione esatta).
2. Asimmetria Composizione-Decomposizione:
   • Il modello mostra invarianza alla complessità nella composizione: una volta apprese le regole base, risolve sequenze di 2-5 passi con la stessa velocità di convergenza.
   • Al contrario, nella decomposizione, l'aumento della lunghezza delle sequenze di supporto ritarda significativamente la convergenza, bloccando il modello in stadi intermedi più a lungo.
3. Finestre di Plasticità Layer-Specifiche: Attraverso interventi di congelamento, gli autori dimostrano che ogni layer del Transformer ha una "finestra di plasticità" temporale specifica. Congelare un layer durante questa finestra ritarda o impedisce il completamento dello stadio successivo, indicando che l'informazione per le fasi avanzate richiede una continua adattabilità di specifici strati.
4. Bias di Adiacenza: Nel compito di scoperta delle pozioni, il modello sfrutta inizialmente la struttura delle ricompense (reward) per indovinare i vicini, mostrando un bias temporaneo che viene poi corretto man mano che apprende la struttura latente completa.

4. Risultati Principali

• Dinamica di Scoperta (Withheld Task): Il modello impara in tre fasi distinte:
  1. Restringe le predizioni all'insieme di supporto ($P[A]$).
  2. Impara a identificare la "metà corretta" del grafo ($P[B|A]$). Curiosamente, l'accuratezza esatta ($P[C|A \cap B]$) sale inizialmente grazie a un bias sulle ricompense adiacenti, prima che il modello impari a distinguere la metà corretta, causando un temporaneo calo dell'accuratezza sulla metà corretta.
  3. Risoluzione finale della corrispondenza esatta.
• Composizione: Le curve di apprendimento per query di lunghezza 2, 3, 4 e 5 sono quasi identiche. Il modello generalizza la composizione delle regole una volta apprese le transizioni atomiche, senza mostrare ritardi legati alla lunghezza della sequenza.
• Decomposizione: Aumentare la lunghezza delle sequenze di supporto ($hlsupport$) ritarda l'acquisizione della fase intermedia ($P[B|A]$, ovvero identificare le pietre raggiungibili). Questo suggerisce che decomporre sequenze lunghe è intrinsecamente più difficile e richiede più tempo per estrarre la struttura latente.
• Plasticità: Il congelamento dei layer trasformatori dopo la prima fase ($P[A]$) causa ritardi significativi nelle fasi successive. Il layer di embedding ha finestre di plasticità molto più brevi rispetto ai layer trasformatori.

5. Significato e Implicazioni

Questo studio offre una visione profonda su come le capacità dei modelli Transformer emergono durante l'addestramento, sfatando l'idea che l'apprendimento sia un processo continuo e omogeneo.

• Interpretabilità: Dimostra che i plateau nelle curve di perdita/accuratezza non sono casuali, ma corrispondono a transizioni di fase nell'acquisizione di abilità specifiche e interpretabili.
• Generalizzazione: Conferma che i Transformer possono imparare strutture latenti vere e proprie e non si limitano a memorizzare pattern superficiali, specialmente in compiti con regole composizionali.
• Limiti Architetturali: Evidenzia una asimmetria fondamentale: i modelli sono eccellenti nel comporre regole note (forward composition), ma faticano a decomporre esempi complessi per risalire alle regole atomiche (inverse decomposition), specialmente quando la complessità dei dati di supporto aumenta.
• Prospettive Future: I risultati suggeriscono che l'ottimizzazione e la stabilità dell'addestramento dipendono criticamente dal mantenimento della plasticità in specifici layer durante finestre temporali precise, offrendo nuove direzioni per migliorare l'efficienza dell'addestramento e la comprensione dei meccanismi di "grokking" (apprendimento improvviso).

In sintesi, il paper fornisce una prova empirica rigorosa che l'apprendimento delle strutture latenti nei Transformer è un processo gerarchico, sequenziale e soggetto a vincoli temporali specifici per ogni strato della rete.