LLM 프리트레이닝에서 Grokking을 어디서 찾을 수 있을까요? 테스트 없이 암기-일반화 과정을 모니터링하기

Grokking, 즉 훈련 손실이 수렴한 후에도 테스트 성능이 계속 개선되는 현상은 최근 신경망 훈련에서 목격되어 일반화 메커니즘과 추론 등의 새로운 능력들이 신비롭게 여겨지고 있습니다. 이전 연구에서는 주로 작은 모델을 수천 에포크 동안 장난감이나 매우 특정한 작업에 대해 훈련시키는 반면, 우리는 70억 파라미터 규모의 대형 언어 모델(LLM), 즉 OLMoE의 일회성 사전훈련 과정 중 체크포인트를 대상으로 Grokking에 대한 최초의 연구를 수행했습니다. 우리는 다양한 벤치마크 작업, 수학적 추론, 코드 생성 및 상식/영역별 지식 검색 작업을 포함하여 훈련 손실을 계산하고 일반화 성능을 평가했습니다.우리의 연구는 처음으로 대규모 기반 모델의 사전훈련 과정에서도 Grokking이 발생한다는 것을 확인하였으며, 다른 데이터가 비동기적으로 Grokking 단계에 진입할 수 있음을 보여주었습니다. 또한, LLM 내부 역학을 조사하여 Grokking의 "일반화의 출현"을 더욱 명확히 밝혔습니다. 구체적으로, 훈련 샘플들의 경로(즉, 레이어 간 전문가 선택)가 Grokking 과정에서 무작위적이며 샘플별로 특화된 형태에서 더 구조적이고 샘플 간 공유 가능한 형태로 진화함을 발견하였습니다. 또한, 손실이 수렴한 상태에서도 샘플 경로의 복잡성이 감소하는 것으로 나타났습니다. 이는 기억에서 일반화로의 전환을 시사하며, 지연된 일반화에 대한 기계적인 설명을 제공합니다.연구에서 우리는 경로 거리와 단일 경로 복잡성을 측정하기 위한 두 가지 새로운 지표를 개발하였습니다. 이들 지표는 다양한 다운스트림 작업에서 일반화 성능 개선을 예측하는 능력을 보여주며, 계산이 효율적이고 간단하며 오직 훈련 데이터에만 의존하므로 사전훈련 과정에서 실용적인 가치를 가집니다. 이를 통해 미세조정(finetuning)이나 테스트 없이도 일반화 성능을 모니터링할 수 있게 되었습니다. 이론적으로는 더 구조적인 경로가 모델 복잡성을 줄이고 일반화 한계를 개선함을 보였습니다.