A Pre-Training Analogue of Grokking in Language Models: Tracing Delayed Grammatical Generalization

기술 요약: 언어 모델에서의 사전 학습 단계 그로킹(Grokking) 아날로그

문제 정의

신경망이 학습이 길어짐에 따라 훈련 데이터를 암기하는 단계에서 벗어나 홀드아웃(held-out) 예시에 일반화되는 현상인 **그로킹(grokking)**은, 명시적인 훈련/검증 분할과 수많은 에포크(epoch)가 존재하는 지도 학습 환경에서 광범위하게 연구되어 왔다(Power et al., 2022). 그러나 이러한 고전적 정식화는 대규모 언어 모델(LLM)의 사전 학습(pre-training) 단계로 직접 전이되지 않는다. 사전 학습에서 모델은 대규모의 라벨 없는 코퍼스(예: C4)에 대해 다음 토큰 예측을 위해 최적화되며, 데이터 반복이 제한적이고 다운스트림 평가 작업에 대한 자연스러운 지도 학습 기반의 훈련/검증 분할이 존재하지 않는다. 결과적으로, 모델이 특정 예시를 먼저 적합(fit)시킨 후, 사전 학습 단계 동안 동일한 작업의 보이지 않는 변형들에 대해 나중에 일반화하는지를 측정할 표준적인 메커니즘이 없다.

방법론

이러한 간극을 메우기 위해, 저자들은 사전 학습 중 지연된 일반화를 연구하기 위한 프록시(proxy) 훈련/검증 분할을 구축하는 **노출 기반 프레임워크(exposure-based framework)**를 제안한다.

1. 노출 기반 프록시 분할:

   • 본 연구는 특정 문법 현상을 대상으로 하는 하나의 허용 가능한 문장과 하나의 허용되지 않는 문장을 포함하는 최소 쌍(minimal pairs) 데이터셋인 BLiMP(Warstadt et al., 2020)를 활용한다.
   • 각 최소 쌍에 대해, 저자들은 **핵심 구절(critical phrase)**을 추출한다: 이는 문법적 대조와 필요한 문맥을 포함하는 가장 작은 연속적인 구간이다.
   • 프록시-훈련(노출됨) 분할: 핵심 구절이 분석된 사전 학습 윈도우 내의 C4 토큰에 그대로(verbatim) 나타나는 예시들.
   • 프록시-검증(노출되지 않음) 분할: 핵심 구절이 사전 학습 윈도우 내에 그대로 나타나지 않는 예시들.
   • 이 분할은 단순히 문법 규칙을 다른 문맥에서 보았는지 여부가 아니라, 구절의 '문자 그대로의 노출(verbatim exposure)' 여부를 기준으로 구분한다.

2. 실험 설정:

   • 모델: 35M, 60M, 130M 파라미터를 가진 LLaMA 스타일의 디코더 전용 모델.
   • 학습: 영어 C4 데이터셋으로 2,000 스텝(약 2억 개의 non-padding 토큰) 동안 학습하며, 100 스텝마다 체크포인트를 생성한다.
   • 평가: BLiMP 최소 쌍에 대해 정확도를 측정한다. 모델은 허용 가능한 문장에 더 높은 확률을 할당할 때 정답으로 간주된다.
   • 지표: **지연된 일반화 시차($\Delta t$)**는 프록시-훈련 분할과 프록시-검증 분할이 80% 정확도 임계값에 도달하는 시점 사이의 학습 스텝 차이로 정의된다.

3. 검증 및 분석:

   • 매칭된 무작위 분할 테스트: 관찰된 시차가 분할 크기나 임의적인 예시 차이로 인한 인위적인 결과가 아님을 보장하기 위해, 저자들은 원래의 그룹 크기를 유지하면서 1,000개의 무작위 분할을 생성한다. 이들은 노출 기반 시차를 이 영가설 분포(null distribution)와 비교한다.
   • 표상 분석: 저자들은 문법적 개념 벡터($v_i = h_{gram} - h_{ungram}$)를 구축하여 문법적 문장과 비문(ungrammatical sentences)의 분리 가능성을 측정하고, 이 벡터들의 **유효 랭크(effective rank)**를 계산하여 표상 공간의 차원을 분석한다.
   • 어텐션 분석: 문법적 의존성을 해결하는 데 필요한 **문맥 토큰(context token)**으로부터 **핵심 토큰(critical token)**으로 향하는 어텐션 가중치를 추적하여, 어텐션 강도와 엔트로피의 변화를 측정한다.

주요 결과

1. 지연된 일반화 관찰: 다섯 가지 문법 현상(예: 주어-동사 일치, Existential There Raising) 전반에 걸쳐, 프록시-훈련(노출된) 분할이 프록시-검증(노출되지 않은) 분할보다 일관되게 높은 정확도(80%)에 먼저 도달한다. 이는 모델이 직접 노출된 예시에는 잘 수행하다가, 구조적으로 유사하지만 문자 그대로 노출되지 않은 예시로 일반화하기 전까지의 지연된 일반화, 즉 사전 학습의 그로킹 아날로그가 존재함을 확인시켜 준다.
2. 시차 검증: 매칭된 무작위 분할 순열 테스트를 통해 관찰된 시차가 통계적으로 유의미하며($p < 0.05$), 분할 크기만으로는 설명될 수 없음을 확인했다. 이 시차는 구절의 문자 그대로의 노출과 구체적으로 연결되어 있다.
3. 개념 벡터의 변화:
   • 분리 가능성: 문법적 개념 벡터의 AUROC는 일반적으로 사전 학습 중에 증가하며, 이는 문법적 문장과 비문 사이의 분리 능력이 향상됨을 나타낸다.
   • 차원: 지연된 일반화 전이 이후, 문법적 개념 벡터의 유효 랭크가 증가한다. 이는 문법적 개념이 몇 개의 지배적인 방향으로 붕괴되지 않고, 더 높은 차원의 부분 공간에 걸쳐 분포하게 됨을 시사한다.
4. 어텐션 재구성: 일반화로의 전이는 국소적인 어텐션 변화를 동반한다. 구체적으로, 핵심 토큰에서 관련 문맥 토큰으로 향하는 어텐션이 소수의 특정 어텐션 헤드에서 증가한다. 이러한 변화는 현상 특이적(예: 주어-동사 일치와 한정사-명사 일치에 대해 서로 다른 레이어/헤드가 작용)이며, 모델 전체의 균일한 변화가 아니다.

주요 기여

1. 프레임워크: 명시적인 다운스트림 감독 없이도 LLM 사전 학습 중 그로킹과 유사한 역학을 연구할 수 있게 하는 노출 기반 프레임워크를 도입하였다.
2. 실증적 증거: 다섯 가지 문법 현상에 대해, 핵심 구절이 문자 그대로 노출된 예시들이 그렇지 않은 예시들보다 유의미하게 더 빨리 높은 정확도에 도달함을 입증하였다.
3. 강건성: 관찰된 시차가 데이터셋 분할에 의한 인위적 결과가 아님을 매칭된 무작위 분할 테스트를 통해 검증하였다.
4. 메커니즘적 통찰: 일반화로의 행동적 전이가 문법 특화적 개념 벡터의 기하학적 변화(차원 증가)와 국소적인 어텐션 재구성과 결합되어 있다는 증거를 제시하였다.

의의 및 주장

본 논문은 "고전적 그로킹 설정의 사전 학습 아날로그"를 제공한다고 주장한다. 코퍼스 노출을 훈련 데이터의 프록시로 활용함으로써, 저자들은 지연된 적응(fitting)에서 일반화로의 전이가 비지도 방식인 다음 토큰 예측 환경의 사전 학습에서도 관찰될 수 있음을 보여준다.

저자들은 이러한 발견이 코퍼스 노출이 그로킹과 유사한 지연된 일반화를 연구하는 데 사용될 수 있음을 시사한다고 강조한다. 메커니즘 분석은 이러한 행동적 전환이 단순한 성능 지표의 변화가 아니라, 모델의 내부 표상에 뿌리를 두고 있음을 밝혀낸다. 즉, 문법적 개념은 더 분리 가능해지고 더 높은 차원의 부분 공간을 점유하며, 어텐션 메커니즘은 의존성 관련 문맥 토큰에 더 강력하게 집중하도록 재구성된다.

본 연구는 "프록시-검증" 분할이 반드시 문법 자체가 완전히 생소한 예시가 아니라, 단지 문자 그대로의 노출이 없는 예시를 의미한다는 점을 언급하며 연구의 범위를 겸허히 제한한다. 또한, 이 프레임워크는 핵심 구절을 문장 내 하위 구간(sub-sentential spans)으로 정의할 수 있는 문법 현상에 국한되어 있다.