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

TL;DR

This study investigates the emergence of grokking in large language model pretraining, revealing how models transition from memorization to generalization and proposing low-cost metrics to monitor this process.

Contribution

It is the first to analyze grokking in practical LLM pretraining, especially in mixture-of-experts models, and introduces data pathway metrics for monitoring generalization.

Findings

Grokking occurs in MoE LLM pretraining with asynchronous local stages.

Training data pathways evolve from random to structured, indicating a memorization-to-generalization transition.

Proposed metrics effectively track model generalization without costly evaluations.

Abstract

This paper presents the first study of grokking in practical LLM pretraining. Specifically, we investigate when an LLM memorizes the training data, when its generalization on downstream tasks starts to improve, and what happens if there is a lag between the two. Unlike existing works studying when a small model generalizes to limited and specified tasks during thousands epochs' training on algorithmic data, we focus on a practical setting for LLMs, i.e., one-epoch pretraining of next-token prediction on a cross-domain, large-scale corpus, and generalization on diverse benchmark tasks covering math/commonsense reasoning, code generation, and domain-specific retrieval. Our study, for the first time, verifies that grokking still emerges in pretraining mixture-of-experts (MoE) LLMs, though different local data groups may enter their grokking stages asynchronously due to the heterogeneity of their distributions and attributions to others. To find a mechanistic interpretation of this local grokking, we investigate the dynamics of training data's pathways (i.e., expert choices across layers in MoE). Our primary discovery is that the pathways evolve from random, non-smooth across layers, instance-specific to more structured and transferable across samples, despite the converged pretraining loss. This depicts a transition from memorization to generalization. Two novel metrics are developed to quantify these patterns: one computes the pathway similarity between samples, while the other measures the consistency of aggregated experts between subsequent layers for each sample. These training data based metrics induce zero cost but can faithfully track and monitor the generalization of LLMs on downstream tasks, which, in conventional settings, requires costly instruction tuning and benchmark evaluation.