Grokking in LLM Pretraining? Monitor Memorization-to-Generalization without Test
Ziyue Li, Chenrui Fan, Tianyi Zhou
TL;DR
The paper addresses grokking in practical LLM pretraining, showing that memorization and delayed generalization occur locally across heterogeneous data groups during one-epoch pretraining of a 7B MoE model. It develops two pathway-based metrics, $D_{\mathrm{path}}$ and $C_i$, to monitor memorization-to-generalization without finetuning and demonstrates strong correlations between these metrics and downstream generalization across math, code, commonsense, and domain-specific benchmarks. The authors provide a theoretical link via a routing kernel $\Theta_{\text{route}}$ and an effective dimension $d_{\rm eff}$, arguing that a collapse of routing complexity under NTK-like assumptions drives the generalization boost after memorization. Practically, the work offers zero-cost, real-time monitoring tools for pretraining, with implications for more transparent, scalable evaluation and guidance on when and how pretraining data contribute to downstream capabilities, while outlining future work to extend these insights to dense architectures.
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.