Dove trovare il Grokking nel Pretraining dei Modelli Linguistici? Monitorare la Transizione da Memorizzazione a Generalizzazione senza Test

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

June 26, 2025

Autori: Ziyue Li, Chenrui Fan, Tianyi Zhou

cs.AI

Abstract

Grokking, ovvero il miglioramento continuo delle prestazioni sui test molto tempo dopo la convergenza della loss di addestramento, è stato recentemente osservato nell'addestramento di reti neurali, rendendo misteriosi i meccanismi di generalizzazione e altre capacità emergenti come il ragionamento. Mentre gli studi precedenti addestrano solitamente modelli di piccole dimensioni su pochi compiti giocattolo o altamente specifici per migliaia di epoche, noi conduciamo il primo studio sul grokking utilizzando checkpoint durante il pretraining in un'unica passata di un modello linguistico di grandi dimensioni (LLM) da 7B, ovvero OLMoE. Calcoliamo la loss di addestramento e valutiamo la generalizzazione su una varietà di benchmark, tra cui ragionamento matematico, generazione di codice e compiti di recupero di conoscenza di senso comune o specifica di dominio. Il nostro studio, per la prima volta, verifica che il grokking si verifica ancora durante il pretraining di modelli di fondazione su larga scala, sebbene diversi dati possano entrare nelle fasi di grokking in modo asincrono. Demistifichiamo ulteriormente l'"emergenza della generalizzazione" del grokking investigando le dinamiche interne degli LLM. In particolare, scopriamo che i percorsi dei campioni di addestramento (ovvero, le scelte degli esperti attraverso i livelli) evolvono da casuali e specifici per l'istanza a più strutturati e condivisibili tra i campioni durante il grokking. Inoltre, la complessità del percorso di un campione si riduce nonostante la loss convergente. Questi risultati indicano una conversione da memorizzazione a generalizzazione, fornendo una spiegazione meccanicistica della generalizzazione ritardata. Nello studio, sviluppiamo due nuove metriche per quantificare la distanza dei percorsi e la complessità di un singolo percorso. Dimostriamo la loro capacità di prevedere il miglioramento della generalizzazione su una varietà di compiti downstream. Sono efficienti, semplici da calcolare e dipendono esclusivamente dai dati di addestramento. Pertanto, hanno un valore pratico per il pretraining, permettendoci di monitorare le prestazioni di generalizzazione senza fine-tuning e test. Teoricamente, dimostriamo che percorsi più strutturati riducono la complessità del modello e migliorano il limite di generalizzazione.

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.