PRISM: 미드 트레이닝에서의 능력 유지 및 상호작용의 원리 규명

대규모 언어 모델의 성능을 극대화하기 위해 사전 학습과 정렬 사이에 수행되는 미드 트레이닝의 효과를 체계적으로 분석했다. 특히 강화 학습의 성공을 위해서는 미드 트레이닝이 필수적인 기초 단계임을 입증하여 효율적인 모델 개발 가이드를 제공한다.

기존 LLM 학습은 대규모 데이터로 기초를 닦는 Pre-training과 인간의 의도에 맞추는 Alignment의 2단계 구조였다. 하지만 최근에는 특정 도메인의 고품질 데이터를 주입하는 Mid-training 단계가 추가되고 있으나, 이 단계가 모델의 내부 가중치와 후속 RL 과정에 어떤 영향을 주는지에 대한 체계적인 연구가 부족했다.

PRISM 연구는 미드 트레이닝이 모델의 가중치 공간을 RL이 효과적으로 작동할 수 있는 최적의 상태로 재배치하는 역할을 수행함을 밝혀냈다. 이는 단순히 지식을 추가하는 수준을 넘어, 모델 파라미터의 90% 이상을 근본적으로 재구성하여 복잡한 추론 사슬을 생성할 수 있는 구조적 토대를 마련하는 과정이다.

결과적으로 미드 트레이닝은 모델이 문제를 단계적으로 분해하여 생각하는 사고 방식을 배우게 하며, 이후의 RL은 이 사고 방식을 더 효율적이고 정확하게 다듬는 역할을 한다. 이러한 단계적 접근은 베이스 모델에 직접 RL을 적용할 때 발생하는 학습 불안정성이나 성능 정체 문제를 해결하는 핵심 열쇠가 된다.

3B에서 24B 규모의 7개 베이스 모델을 대상으로 약 27B 토큰의 고품질 데이터 믹스를 사용하여 미드 트레이닝을 수행했다. 데이터는 일반 웹 데이터(DCLM-EDU), 수학(Open-R1, Nemotron), 코드(StarCoder2), 과학(OpenThoughts3) 및 채팅 데이터를 정교하게 혼합하여 구성했다.

미드 트레이닝 후 GRPO 알고리즘을 사용하여 강화 학습을 진행했다. [모델이 생성한 여러 응답의 보상을 계산] → [평균 보상 대비 상대적 이득을 산출] → [이득이 큰 응답의 확률을 높이는 방향으로 업데이트] 과정을 통해 추론 성능을 최적화했다.

모델의 내부 변화를 측정하기 위해 CKA와 가중치 발산 분석을 도입했다. [두 체크포인트의 은닉 상태 행렬 X, Y를 입력으로] → [커널 행렬 간의 정렬도를 계산] → [0에서 1 사이의 유사도 점수를 도출] 하여 표현 공간의 기하학적 변화를 정량화했다.

미드 트레이닝을 거친 Granite-3.3-8B 모델은 MATH500 벤치마크에서 베이스 모델(26.09%) 대비 81.11%로 급격한 성능 향상을 보였다. AIME24 점수 또한 0.46에서 37.18로 상승하며 고난도 수학 문제 해결 능력이 크게 개선됐다.

전체 PRISM → RL 파이프라인을 적용했을 때, 6개 추론 벤치마크의 가중 평균 점수는 12점 미만에서 2942점으로 34배 향상됐다. 반면 미드 트레이닝 없이 RL만 적용한 경우 AIME 점수는 거의 0에 머물렀다.

가중치 분석 결과 미드 트레이닝은 파라미터의 90% 이상을 변경하는 조밀한 재구성을 수행하는 반면, RL은 단 5% 미만의 파라미터만 변경하는 희소한 미세 조정을 수행하면서도 성능을 추가로 8~12점 더 끌어올렸다.

미드 트레이닝은 모델 파라미터의 전역적인 재구성을 유도하며, 특히 MLP 레이어에서 가장 큰 가중치 변화가 발생한다. 이는 모델이 새로운 도메인 지식을 내부 표현 공간에 조밀하게 통합하는 과정임을 시사한다.

RL 단계의 가중치 변화는 매우 희소하며, 주로 Attention 레이어의 Value(V) 및 Output(O) 프로젝션 행렬에 집중된다. Mamba 아키텍처의 경우 SSM 파라미터는 거의 변하지 않고 유지되는 특성을 보인다.

CKA 분석을 통해 RL이 미드 트레이닝으로 형성된 표현 공간의 기하학적 구조를 0.998 이상의 높은 유사도로 보존함을 확인했다. 이는 RL이 새로운 표현을 만들기보다 기존 표현을 활용하여 최적의 응답 경로를 찾는 과정임을 의미한다.

미드 트레이닝 시 8k 컨텍스트 윈도우를 사용하면 사전 학습된 모델의 장기 문맥 능력이 저하되나, 15%의 베이스 모델 가중치를 병합하고 짧은 확장 학습을 추가함으로써 성능을 효과적으로 복구할 수 있다.

RL 데이터 선택 시 단일 모델(Granite-3.3-8B)로 필터링한 데이터를 모든 모델에 동일하게 적용했으므로, 각 모델에 최적화된 데이터 큐레이션을 적용할 경우 더 높은 성능 향상이 가능할 수 있다. 또한 24B 이상의 초대형 모델에 대한 검증은 향후 과제로 남아 있다.