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

Technische Samenvatting: Een Pre-training Analoog van Grokking in Taalmodellen

Probleemstelling

Het fenomeen van grokking—waarbij neurale netwerken overgaan van het memoriseren van trainingsdata naar het generaliseren naar niet-geziene voorbeelden na een langdurige trainingsperiode (Power et al., 2022)—is uitgebreid bestudeerd in gesuperviseerde settings met expliciete train/validatie-splits en vele epochs. Echter, deze klassieke formulering vertaalt zich niet direct naar de pre-training van Large Language Models (LLM's). Tijdens pre-training worden modellen geoptimaliseerd voor next-token prediction over een grote, ongelabelde corpus (bijv. C4) met beperkte datarepetitie en zonder natuurlijke gesuperviseerde train/validatie-split over downstream evaluatietaken. Bijgevolg is er geen standaardmechanisme om te meten of een model eerst specifieke voorbeelden past en pas later generaliseert naar ongeziene variaties van dezelfde taak tijdens de pre-training fase.

Methodologie

Om dit gat te dichten, stellen de auteurs een exposure-gebaseerd framework voor om een proxy train/validatie-split te construeren voor het bestuderen van vertraagde generalisatie tijdens pre-training.

1. Exposure-gebaseerde Proxy-split:

   • De studie maakt gebruik van BLiMP (Warstadt et al., 2020), een dataset van minimale paren bestaande uit één acceptabele en één onacceptabele zin die gericht zijn op specifieke grammaticale fenomenen.
   • Voor elk minimaal paar extraheren de auteurs een kritieke frase: de kleinste continue span die het grammaticale contrast en de noodzakelijke context bevat.
   • Proxy-Train (Exposed) Split: Voorbeelden waar de kritieke frase letterlijk voorkomt in de C4-tokens die gezien zijn tijdens de geanalyseerde pre-training window.
   • Proxy-Validation (Unexposed) Split: Voorbeelden waar de kritieke frase niet letterlijk voorkomt in de pre-training window.
   • Deze split onderscheidt op basis van letterlijke exposure, niet of het model de onderliggende grammaticale regel in andere contexten heeft gezien.

2. Experimentele Opzet:

   • Modellen: LLaMA-stijl decoder-only modellen met 35M, 60M en 130M parameters.
   • Training: Getraind op Engelse C4 voor 2.000 stappen (ongeveer 200M non-padding tokens) met checkpoints elke 100 stappen.
   • Evaluatie: Nauwkeurigheid wordt gemeten op BLiMP minimale paren. Het model is correct als het een hogere waarschijnlijkheid toekent aan de acceptabele zin.
   • Metriek: De Delayed-Generalization Lag ($\Delta t$) wordt gedefinieerd als het verschil in trainingsstappen tussen wanneer de proxy-train split en de proxy-validation split een nauwkeurigheidsdrempel van 80% bereiken.

3. Validatie en Analyse:

   • Matched Random-Split Test: Om te waarborgen dat de geobserveerde lag geen artefact is van de split-grootte of willekeurige verschillen tussen voorbeelden, genereren de auteurs 1.000 willekeurige splits die de oorspronkelijke groepsgroottes behouden. Ze vergelijken de exposure-gebaseerde lag met deze nulverdeling.
   • Representatie-analyse: Zij construeren grammaticale conceptvectoren ($v_i = h_{gram} - h_{ungram}$) om de scheidbaarheid van grammaticale versus ongrammaticale zinnen te meten en berekenen de effectieve rang van deze vectoren om de dimensionaliteit van de representatieruimte te analyseren.
   • Attention-analyse: Zij volgen de attention weights van de kritieke token naar de context token die nodig is om de grammaticale afhankelijkheid op te lossen, waarbij veranderingen in attention sterkte en entropie worden gemeten.

Belangrijkste Resultaten

1. Observatie van Vertraagde Generalisatie: Over vijf grammaticale fenomenen (bijv. Subject-Verb Agreement, Existential There Raising) bereikt de proxy-train (exposed) split consistent een hoge nauwkeurigheid (80%) vóór de proxy-validation (unexposed) split. Dit bevestigt een pre-training analoog van grokking waarbij modellen goed presteren op direct geëxposeerde voorbeelden voordat ze generaliseren naar structureel gelijkaardige maar letterlijk niet-geëxposeerde voorbeelden.
2. Validatie van de Lag: Matched random-split permutatietesten bevestigen dat deze lag statistisch significant is ($p < 0.05$) en niet alleen verklaard kan worden door de split-grootte. De lag is specifiek gekoppeld aan de letterlijke kritieke-frase exposure.
3. Veranderingen in Conceptvectoren:
   • Scheidbaarheid: De AUROC van grammaticale conceptvectoren neemt over het algemeen toe tijdens pre-training, wat wijst op een verbeterde scheidbaarheid tussen grammaticale en ongrammaticale zinnen.
   • Dimensionaliteit: Na de delayed-generalization transitie neemt de effectieve rang van grammaticale conceptvectoren toe. Dit suggereert dat grammaticale concepten verdeeld worden over een hoger-dimensionale subspace in plaats van te kollapsen naar enkele dominante richtingen.
4. Attention Reorganisatie: De transitie naar generalisatie gaat gepaard met gelokaliseerde veranderingen in attention. Specifiek neemt de attention van de kritieke token naar de relevante context token toe in een klein aantal specifieke attention heads. Deze veranderingen zijn fenomeen-specifiek (bijv. verschillende lagen/heads voor Subject-Verb Agreement versus Determiner-Noun Agreement) in plaats van een uniforme verschuiving over het hele model.

Belangrijkste Bijdragen

1. Framework: Introductie van een exposure-gebaseerd framework dat het bestuderen van grokking-achtige dynamiek tijdens LLM pre-training mogelijk maakt zonder expliciete downstream supervisie.
2. Empirisch Bewijs: Demonstratie dat over vijf grammaticale fenomenen, voorbeelden met letterlijke kritieke-frase exposure aanzienlijk eerder een hoge nauwkeurigheid bereiken dan voorbeelden zonder dergelijke exposure.
3. Robuustheid: Validatie dat de geobserveerde lag geen artefact is van dataset-splitting via matched random-split tests.
4. Mechanistisch Inzicht: Bewijs dat de gedragsmatige transitie naar generalisatie gepaard gaat met specifieke veranderingen in de geometrie van grammatica-specifieke conceptvectoren (toenemende dimensionaliteit) en gelokaliseerde attention reorganisatie.

Betekenis en Claims

Het artikel claimt een "pre-training analoog van de klassieke grokking setup" te bieden. Door corpus exposure te gebruiken als proxy voor trainingsdata, demonstreren de auteurs dat de vertraagde transitie van fitting naar generalisatie observeerbaar is, zelfs in de ongesuperviseerde, next-token prediction setting van pre-training.

De auteurs benadrukken dat hun bevindingen suggereren dat corpus exposure gebruikt kan worden om grokking-achtige vertraagde generalisatie te bestuderen. De mechanistische analyses onthullen dat deze gedragsmatige verschuiving niet louter een verandering in een prestatie-metriek is, maar geworteld is in de interne representaties van het model: grammaticale concepten worden scheidbaarder en bezetten hoger-dimensionale subspaces, terwijl attention-mechanismen reorganiseren om sterker te focussen op dependency-relevante context tokens.

De studie blijft bescheiden in haar reikwijdte, waarbij wordt opgemerkt dat de "proxy-validation" split voorbeelden vertegenwoordigt zonder letterlijke exposure, maar niet noodzakelijkerwijs voorbeelden met een geheel onbekende grammatica. Het framework is momenteel beperkt tot grammaticale fenomenen waarbij kritieke frases gedefinieerd kunnen worden als sub-sententiële spans.