Wo findet man Grokking im Pretraining von LLM? Überwachung des Übergangs von Auswendiglernen zu Generalisierung ohne Test

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

June 26, 2025

Autoren: Ziyue Li, Chenrui Fan, Tianyi Zhou

cs.AI

Zusammenfassung

Grokking, d.h. die kontinuierliche Verbesserung der Testleistung, lange nachdem der Trainingsverlust konvergiert ist, wurde kürzlich im Training von neuronalen Netzwerken beobachtet, was den Mechanismus der Generalisierung und andere aufkommende Fähigkeiten wie das logische Denken rätselhaft erscheinen lässt. Während frühere Studien in der Regel kleine Modelle auf wenigen einfachen oder hochspezifischen Aufgaben über Tausende von Epochen trainierten, führen wir die erste Studie zu Grokking an Checkpoints während des Einmal-Vortrainings eines großen Sprachmodells (LLM) mit 7B Parametern, nämlich OLMoE, durch. Wir berechnen den Trainingsverlust und bewerten die Generalisierung auf verschiedenen Benchmark-Aufgaben, darunter mathematisches Denken, Code-Generierung und Aufgaben zur Abfrage von Allgemeinwissen bzw. domänenspezifischem Wissen. Unsere Studie bestätigt erstmals, dass Grokking auch beim Vortraining von großskaligen Basismodellen auftritt, obwohl verschiedene Daten asynchron in die Grokking-Phasen eintreten können. Wir entmystifizieren weiterhin das „Aufkommen der Generalisierung“ beim Grokking, indem wir die internen Dynamiken des LLM untersuchen. Insbesondere stellen wir fest, dass sich die Pfade der Trainingsdaten (d.h. die Expertenauswahl über die Schichten hinweg) während des Grokkings von zufälligen, instanzspezifischen zu strukturierteren und zwischen den Proben teilbaren Pfaden entwickeln. Zudem verringert sich die Komplexität eines Probenpfads trotz des konvergierten Verlusts. Dies deutet auf einen Übergang von der Auswendiglernphase zur Generalisierung hin und liefert eine mechanistische Erklärung für die verzögerte Generalisierung. In der Studie entwickeln wir zwei neuartige Metriken, um die Pfaddistanz und die Komplexität eines einzelnen Pfads zu quantifizieren. Wir zeigen deren Fähigkeit, die Verbesserung der Generalisierung auf verschiedenen Downstream-Aufgaben vorherzusagen. Sie sind effizient, einfach zu berechnen und ausschließlich von den Trainingsdaten abhängig. Daher haben sie praktischen Wert für das Vortraining, da sie es uns ermöglichen, die Generalisierungsleistung ohne Feinabstimmung und Tests zu überwachen. Theoretisch zeigen wir, dass strukturiertere Pfade die Modellkomplexität reduzieren und die Generalisierungsgrenze verbessern.

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.