LLM 사전 학습에서 그로킹(Grokking)을 찾는 방법: 테스트 없이 암기에서 일반화로의 전환 모니터링

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

June 26, 2025

저자: Ziyue Li, Chenrui Fan, Tianyi Zhou

cs.AI

초록

그로킹(Grokking), 즉 훈련 손실이 수렴한 후에도 오랜 시간 동안 테스트 성능이 계속 향상되는 현상은 최근 신경망 훈련에서 관찰되며, 일반화 메커니즘과 추론과 같은 새롭게 나타나는 능력들을 이해하기 어렵게 만들고 있다. 기존 연구들은 주로 소규모 모델을 몇 가지 간단하거나 매우 특수한 작업에 대해 수천 에포크 동안 훈련시키는 반면, 본 연구에서는 7B 규모의 대형 언어 모델(LLM), 즉 OLMoE의 원패스(one-pass) 사전 훈련 중 체크포인트에서 그로킹 현상을 처음으로 연구한다. 우리는 훈련 손실을 계산하고 수학적 추론, 코드 생성, 상식 및 도메인 특화 지식 검색 작업을 포함한 다양한 벤치마크 작업에서 일반화 성능을 평가한다. 본 연구는 대규모 기반 모델의 사전 훈련에서도 그로킹이 여전히 발생함을 처음으로 확인하며, 서로 다른 데이터가 비동기적으로 그로킹 단계에 진입할 수 있음을 보여준다. 또한, 우리는 LLM의 내부 동역학을 조사함으로써 그로킹의 "일반화의 출현"을 해명한다. 구체적으로, 그로킹 동안 훈련 샘플의 경로(즉, 계층 간 전문가 선택)가 무작위적이고 인스턴스 특정적에서 더 구조화되고 샘플 간에 공유 가능한 형태로 진화함을 발견한다. 또한, 손실이 수렴한 후에도 샘플 경로의 복잡도가 감소한다. 이러한 현상들은 암기에서 일반화로의 전환을 나타내며, 지연된 일반화에 대한 기계적 설명을 제공한다. 본 연구에서 우리는 경로 거리와 단일 경로의 복잡도를 정량화하기 위한 두 가지 새로운 지표를 개발한다. 이 지표들은 다양한 다운스트림 작업에서 일반화 성능 향상을 예측할 수 있는 능력을 보여준다. 이들은 효율적이며 계산이 간단하고 훈련 데이터에만 의존한다. 따라서 사전 훈련에서 실용적인 가치를 가지며, 파인튜닝과 테스트 없이도 일반화 성능을 모니터링할 수 있게 해준다. 이론적으로, 더 구조화된 경로는 모델 복잡도를 줄이고 일반화 경계를 개선함을 보여준다.

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.