A Pre-Training Analogue of Grokking in Language Models: Tracing Delayed Grammatical Generalization

Sintesi Tecnica: Un Analogo del Grokking nel Pre-training dei Modelli Linguistici

Definizione del Problema

Il fenomeno del grokking — ovvero il momento in cui le reti neurali passano dalla memorizzazione dei dati di addestramento alla generalizzazione su esempi non visti durante un prolungato periodo di addestramento (Power et al., 2022) — è stato ampiamente studiato in contesti supervisionati con espliciti split train/validation e numerosi epoch. Tuttavia, questa formulazione classica non si traduce direttamente nel pre-training dei Large Language Models (LLM). Nel pre-training, i modelli sono ottimizzati per la predizione del next-token su un vasto corpus non etichettato (ad esempio, C4) con una limitata ripetizione dei dati e senza uno split naturale supervisionato train/validation per i task di valutazione a valle. Di conseguenza, non esiste un meccanismo standard per misurare se un modello prima si adatti a esempi specifici e solo successivamente generalizzi a variazioni inedite dello stesso task durante la fase di pre-training.

Metodologia

Per colmare questa lacuna, gli autori propongono un framework basato sull'esposizione per costruire uno split proxy train/validation per studiare la generalizzazione ritardata durante il pre-training.

1. Split Proxy Basato sull'Esposizione:

   • Lo studio utilizza BLiMP (Warstadt et al., 2020), un dataset di coppie minimali contenenti una frase accettabile e una inaccettabile che mirano a specifici fenomeni grammaticali.
   • Per ogni coppia minimale, gli autori estraggono una frase critica: il più piccolo intervallo continuo che contiene il contrasto grammaticale e il contesto necessario.
   • Split Proxy-Train (Esposto): Esempi in cui la frase critica appare verbatim nei token di C4 osservati durante la finestra di analisi del pre-training.
   • Split Proxy-Validation (Non Esposto): Esempi in cui la frase critica non appare verbatim nella finestra di pre-training.
   • Questo split distingue in base all'esposizione verbatim, non in base al fatto che il modello abbia visto la regola grammaticale sottostante in altri contesti.

2. Configurazione Sperimentale:

   • Modelli: Modelli decoder-only di tipo LLaMA con 35M, 60M e 130M di parametri.
   • Addestramento: Addestrati su C4 inglese per 2.000 step (circa 200M di token non-padding) con checkpoint ogni 100 step.
   • Valutazione: L'accuratezza è misurata sulle coppie minimali di BLiMP. Il modello è considerato corretto se assegna una probabilità maggiore alla frase accettabile.
   • Metrica: Il Ritardo di Generalizzazione Ritardata ($\Delta t$) è definito come la differenza in step di addestramento tra quando lo split proxy-train e lo split proxy-validation raggiungono una soglia di accuratezza dell'80%.

3. Validazione e Analisi:

   • Test di Split Casuale Corrispondente: Per garantire che il ritardo osservato non sia un artefatto della dimensione dello split o di differenze arbitrarie tra gli esempi, gli autori generano 1.000 split casuali preservando le dimensioni originali dei gruppi. Confrontano il ritardo basato sull'esposizione con questa distribuzione nulla.
   • Analisi delle Rappresentazioni: Costruiscono vettori di concetti grammaticali ($v_i = h_{gram} - h_{ungram}$) per misurare la separabilità tra frasi grammaticali e ungrammaticali e calcolano il rango effettivo di questi vettori per analizzare la dimensionalità dello spazio di rappresentazione.
   • Analisi dell'Attenzione: Monitorano i pesi dell'attenzione dal token critico al token di contesto necessario per risolvere la dipendenza grammaticale, misurando i cambiamenti nella forza e nell'entropia dell'attenzione.

Risultati Chiave

1. Osservazione della Generalizzazione Ritardata: Attraverso cinque fenomeni grammaticali (ad esempio, Concordanza Soggetto-Verbo, Existential There Raising), lo split proxy-train (esposto) raggiunge costantemente un'alta accuratezza (80%) prima dello split proxy-validation (non esposto). Ciò conferma un analogo del grokking nel pre-training, dove i modelli performano bene sugli esempi direttamente esposti prima di generalizzare a esempi strutturalmente simili ma non esposti in forma verbatim.
2. Validazione del Ritardo: I test di permutazione con split casuali corrispondenti confermano che questo ritardo è statisticamente significativo ($p < 0.05$) e non può essere spiegato solo dalla dimensione dello split. Il ritardo è specificamente legato all'esposizione verbatim della frase critica.
3. Cambiamenti nei Vettori di Concetto:
   • Separabilità: L'AUROC dei vettori di concetti grammaticali aumenta generalmente durante il pre-training, indicando una migliore separabilità tra frasi grammaticali e ungrammaticali.
   • Dimensionalità: Dopo la transizione alla generalizzazione ritardata, il rango effettivo dei vettori di concetti grammaticali aumenta. Ciò suggerisce che i concetti grammaticali diventano distribuiti in un sottospazio a dimensionalità più elevata piuttosto che collassare in poche direzioni dominanti.
4. Riorganizzazione dell'Attenzione: La transizione verso la generalizzazione è accompagnata da cambiamenti localizzati nell'attenzione. Nello specifico, l'attenzione dal token critico al token di contesto rilevante aumenta in un piccolo numero di teste di attenzione specifiche. Questi cambiamenti sono fenomeni-specifici (ad esempio, layer/teste diverse per la Concordanza Soggetto-Verbo rispetto alla Concordanza Determiner-Noun) piuttosto che un cambiamento uniforme in tutto il modello.

Contributi Chiave

1. Framework: Introduzione di un framework basato sull'esposizione che consente lo studio delle dinamiche simili al grokking durante il pre-training degli LLM senza esplicita supervisione a valle.
2. Evidenza Empirica: Dimostrazione che, attraverso cinque fenomeni grammaticali, gli esempi con esposizione verbatim della frase critica raggiungono un'alta accuratezza significativamente prima di quelli senza tale esposizione.
3. Robustezza: Validazione del fatto che il ritardo osservato non è un artefatto della divisione del dataset tramite test di split casuale corrispondente.
4. Insight Meccanicistico: Evidenza che la transizione comportamentale verso la generalizzazione è accompagnata da cambiamenti specifici nella geometria dei vettori di concetti grammaticali (aumento della dimensionalità) e da una riorganizzazione localizzata dell'attenzione.

Significato e Rivendicazioni

L'articolo sostiene di fornire un "analogo del pre-training del setup classico del grokking". Sfruttando l'esposizione al corpus come proxy per i dati di addestramento, gli autori dimostrano che la transizione ritardata dal fitting alla generalizzazione è osservabile anche nell'impostazione di next-token prediction non supervisionata del pre-training.

Gli autori sottolineano che i loro risultati suggeriscono che l'esposizione al corpus può essere utilizzata per studiare la generalizzazione ritardata simile al grokking. L'analisi meccanicistica rivela che questo spostamento comportamentale non è meramente un cambiamento della metrica di performance, ma è radicato nelle rappresentazioni interne del modello: i concetti grammaticali diventano più separabili e occupano sottospazi a dimensionalità più elevata, mentre i meccanismi di attenzione si riorganizzano per focalizzarsi più fortemente sui token di contesto rilevanti per la dipendenza.

Lo studio rimane di portata modesta, osservando che lo "split proxy-validation" rappresenta esempi senza esposizione verbatim, non necessariamente esempi con una grammatica interamente inedita. Il framework è attualmente limitato ai fenomeni grammaticali in cui le frasi critiche possono essere definite come intervalli sub-sentenziali.