Où trouver le Grokking dans le prétraitement des LLM ? Surveiller la transition de la mémorisation à la généralisation sans 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

Résumé

Le phénomène de "grokking", c'est-à-dire l'amélioration continue des performances en test bien après la convergence de la perte d'entraînement, a récemment été observé dans l'entraînement des réseaux de neurones, rendant mystérieux le mécanisme de généralisation ainsi que d'autres capacités émergentes telles que le raisonnement. Alors que les études antérieures entraînaient généralement de petits modèles sur quelques tâches simplistes ou hautement spécifiques pendant des milliers d'époques, nous menons la première étude sur le grokking à partir de points de contrôle (checkpoints) lors de l'entraînement en une passe d'un grand modèle de langage (LLM) de 7 milliards de paramètres, à savoir OLMoE. Nous calculons la perte d'entraînement et évaluons la généralisation sur diverses tâches de référence, incluant le raisonnement mathématique, la génération de code et la récupération de connaissances de bon sens ou spécifiques à un domaine. Notre étude vérifie, pour la première fois, que le grokking se produit également lors de l'entraînement préalable de modèles de fondation à grande échelle, bien que différentes données puissent entrer dans les phases de grokking de manière asynchrone. Nous démystifions en outre l'"émergence de la généralisation" associée au grokking en étudiant la dynamique interne des LLM. Plus précisément, nous constatons que les parcours des échantillons d'entraînement (c'est-à-dire les choix d'experts à travers les couches) évoluent d'un état aléatoire et spécifique à l'instance vers un état plus structuré et partageable entre les échantillons pendant le grokking. De plus, la complexité du parcours d'un échantillon diminue malgré la convergence de la perte. Ces observations indiquent une conversion de la mémorisation vers la généralisation, fournissant une explication mécaniste de la généralisation retardée. Dans cette étude, nous développons deux nouvelles métriques pour quantifier la distance entre les parcours et la complexité d'un parcours individuel. Nous démontrons leur capacité à prédire l'amélioration de la généralisation sur diverses tâches en aval. Ces métriques sont efficaces, simples à calculer et dépendent uniquement des données d'entraînement. Elles ont donc une valeur pratique pour l'entraînement préalable, nous permettant de surveiller les performances de généralisation sans ajustement fin ni tests supplémentaires. Théoriquement, nous montrons que des parcours plus structurés réduisent la complexité du modèle et améliorent la borne de généralisation.

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.