¿Dónde encontrar la comprensión en el preentrenamiento de modelos de lenguaje? Monitorear la transición de memorización a generalización sin pruebas

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

June 26, 2025

Autores: Ziyue Li, Chenrui Fan, Tianyi Zhou

cs.AI

Resumen

Grokking, es decir, el rendimiento en pruebas sigue mejorando mucho después de que la pérdida de entrenamiento haya convergido, ha sido observado recientemente en el entrenamiento de redes neuronales, lo que hace que el mecanismo de generalización y otras capacidades emergentes, como el razonamiento, resulten misteriosos. Mientras que estudios previos suelen entrenar modelos pequeños en unas pocas tareas específicas o de juguete durante miles de épocas, nosotros realizamos el primer estudio de grokking en puntos de control durante el preentrenamiento de un solo paso de un modelo de lenguaje grande (LLM) de 7B, es decir, OLMoE. Calculamos la pérdida de entrenamiento y evaluamos la generalización en diversas tareas de referencia, incluyendo razonamiento matemático, generación de código y tareas de recuperación de conocimiento de sentido común o específico de un dominio. Nuestro estudio, por primera vez, verifica que el grokking sigue ocurriendo durante el preentrenamiento de modelos fundacionales a gran escala, aunque diferentes datos pueden entrar en etapas de grokking de manera asincrónica. Además, desmitificamos la "emergencia de la generalización" del grokking al investigar la dinámica interna de los LLM. Específicamente, encontramos que las rutas de las muestras de entrenamiento (es decir, las elecciones de expertos a través de las capas) evolucionan desde ser aleatorias y específicas de la instancia hasta volverse más estructuradas y compartibles entre muestras durante el grokking. Además, la complejidad de la ruta de una muestra se reduce a pesar de la pérdida convergida. Esto indica una conversión de memorización a generalización, proporcionando una explicación mecanicista de la generalización retrasada. En el estudio, desarrollamos dos métricas novedosas para cuantificar la distancia de las rutas y la complejidad de una sola ruta. Demostramos su capacidad para predecir la mejora en la generalización en diversas tareas posteriores. Estas métricas son eficientes, fáciles de calcular y dependen únicamente de los datos de entrenamiento. Por lo tanto, tienen un valor práctico para el preentrenamiento, permitiéndonos monitorear el rendimiento de generalización sin necesidad de ajuste fino ni pruebas. Teóricamente, demostramos que las rutas más estructuradas reducen la complejidad del modelo y mejoran el límite de generalización.

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.