Waar vind je Grokking in LLM-pretraining? Monitor Memorisatie-naar-Generalizatie zonder Test

Where to find Grokking in LLM Pretraining? Monitor Memorization-to-Generalization without Test

June 26, 2025

Auteurs: Ziyue Li, Chenrui Fan, Tianyi Zhou

cs.AI

Samenvatting

Grokking, d.w.z. dat de testprestaties blijven verbeteren lang nadat de trainingsverlies is geconvergeerd, is recentelijk waargenomen bij het trainen van neurale netwerken, wat het mechanisme van generalisatie en andere opkomende capaciteiten zoals redeneren mysterieus maakt. Terwijl eerdere studies meestal kleine modellen trainen op een paar speelse of zeer specifieke taken voor duizenden epochs, voeren wij de eerste studie uit van grokking op checkpoints tijdens een eenmalige pretraining van een groot taalmodel (LLM) van 7B, namelijk OLMoE. We berekenen het trainingsverlies en evalueren de generalisatie op diverse benchmarktaken, waaronder wiskundig redeneren, codegeneratie en taken voor het ophalen van algemene kennis en domeinspecifieke kennis. Onze studie bevestigt voor het eerst dat grokking nog steeds plaatsvindt tijdens de pretraining van grootschalige foundationmodellen, hoewel verschillende gegevens asynchroon de grokkingfasen kunnen ingaan. We ontrafelen verder het "opkomen van generalisatie" van grokking door de interne dynamiek van LLM's te onderzoeken. Specifiek vinden we dat de paden van trainingsvoorbeelden (d.w.z. expertkeuzes over lagen heen) evolueren van willekeurig, voorbeeld-specifiek naar meer gestructureerd en deelbaar tussen voorbeelden tijdens grokking. Ook neemt de complexiteit van het pad van een voorbeeld af ondanks het geconvergeerde verlies. Dit wijst op een conversie van memorisatie naar generalisatie, wat een mechanistische verklaring biedt voor vertraagde generalisatie. In de studie ontwikkelen we twee nieuwe metrieken om de padafstand en de complexiteit van een enkel pad te kwantificeren. We tonen hun vermogen aan om de verbetering van generalisatie op diverse downstreamtaken te voorspellen. Ze zijn efficiënt, eenvoudig te berekenen en uitsluitend afhankelijk van trainingsgegevens. Daarom hebben ze praktische waarde voor pretraining, waardoor we de generalisatieprestaties kunnen monitoren zonder finetuning en testen. Theoretisch tonen we aan dat meer gestructureerde paden de modelcomplexiteit verminderen en de generalisatiegrens verbeteren.

English

Grokking, i.e., test performance keeps improving long after training loss converged, has been recently witnessed in neural network training, making the mechanism of generalization and other emerging capabilities such as reasoning mysterious. While prior studies usually train small models on a few toy or highly-specific tasks for thousands of epochs, we conduct the first study of grokking on checkpoints during one-pass pretraining of a 7B large language model (LLM), i.e., OLMoE. We compute the training loss and evaluate generalization on diverse benchmark tasks, including math reasoning, code generation, and commonsense/domain-specific knowledge retrieval tasks. Our study, for the first time, verifies that grokking still happens in the pretraining of large-scale foundation models, though different data may enter grokking stages asynchronously. We further demystify grokking's "emergence of generalization" by investigating LLM internal dynamics. Specifically, we find that training samples' pathways (i.e., expert choices across layers) evolve from random, instance-specific to more structured and shareable between samples during grokking. Also, the complexity of a sample's pathway reduces despite the converged loss. These indicate a memorization-to-generalization conversion, providing a mechanistic explanation of delayed generalization. In the study, we develop two novel metrics to quantify pathway distance and the complexity of a single pathway. We show their ability to predict the generalization improvement on diverse downstream tasks. They are efficient, simple to compute and solely dependent on training data. Hence, they have practical value for pretraining, enabling us to monitor the generalization performance without finetuning and test. Theoretically, we show that more structured pathways reduce model complexity and improve the generalization bound.