Onde encontrar o Grokking no Pré-treinamento de LLM? Monitorando a Transição de Memorização para Generalização sem Testes

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

Resumo

Grokking, ou seja, o desempenho em testes continua melhorando muito tempo após a convergência da perda de treinamento, foi recentemente observado no treinamento de redes neurais, tornando o mecanismo de generalização e outras capacidades emergentes, como o raciocínio, misteriosos. Enquanto estudos anteriores geralmente treinam modelos pequenos em algumas tarefas específicas ou de brinquedo por milhares de épocas, realizamos o primeiro estudo de grokking em checkpoints durante o pré-treinamento de uma única passagem de um modelo de linguagem grande (LLM) de 7B, ou seja, o OLMoE. Calculamos a perda de treinamento e avaliamos a generalização em diversas tarefas de benchmark, incluindo raciocínio matemático, geração de código e tarefas de recuperação de conhecimento de senso comum/específico de domínio. Nosso estudo, pela primeira vez, verifica que o grokking ainda ocorre no pré-treinamento de modelos de base em grande escala, embora diferentes dados possam entrar em estágios de grokking de forma assíncrona. Além disso, desmistificamos a "emergência da generalização" do grokking investigando a dinâmica interna do LLM. Especificamente, descobrimos que os caminhos das amostras de treinamento (ou seja, as escolhas de especialistas entre as camadas) evoluem de aleatórios e específicos para a instância para mais estruturados e compartilháveis entre as amostras durante o grokking. Além disso, a complexidade do caminho de uma amostra diminui apesar da perda convergida. Isso indica uma conversão de memorização para generalização, fornecendo uma explicação mecanicista da generalização tardia. No estudo, desenvolvemos duas novas métricas para quantificar a distância do caminho e a complexidade de um único caminho. Mostramos sua capacidade de prever a melhoria da generalização em diversas tarefas downstream. Elas são eficientes, simples de calcular e dependem exclusivamente dos dados de treinamento. Portanto, têm valor prático para o pré-treinamento, permitindo-nos monitorar o desempenho de generalização sem ajuste fino e teste. Teoricamente, mostramos que caminhos mais estruturados reduzem a complexidade do modelo e melhoram o limite de generalização.

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.