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

왜 중요한가

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

핵심 기여

미드 트레이닝을 통한 추론 능력의 비약적 향상

약 270억 개의 고품질 토큰을 사용한 미드 트레이닝만으로 수학(+15~40점), 코드(+5~12점), 과학(+6~13점) 벤치마크에서 일관된 성능 향상을 달성했다.

강화 학습(RL) 효과의 극대화 및 필수 조건 확인

미드 트레이닝을 거친 모델은 RL 적용 시 추론 성능이 3~4배 개선되지만, 미드 트레이닝 없이 RL을 직접 적용한 베이스 모델은 대부분 성능 개선에 실패하거나 오히려 하락함을 확인했다.

미드 트레이닝과 RL의 서로 다른 작동 메커니즘 규명

가중치 분석 결과, 미드 트레이닝은 모델 파라미터의 90% 이상을 조밀하게 재구성하는 반면, RL은 약 5%의 파라미터만을 희소하게 미세 조정하며 기존의 표현 기하학적 구조를 유지한다.

데이터 구성의 중요성 및 장기 문맥 복구 전략 제시

데이터 믹스의 영향은 RL 단계보다 미드 트레이닝 단계에서 훨씬 크며, 미드 트레이닝 중 손실된 장기 문맥 처리 능력은 짧은 확장 학습과 모델 병합을 통해 복구 가능하다.

핵심 아이디어 이해하기

기존 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점 미만에서 29~~42점으로 3~~4배 향상됐다. 반면 미드 트레이닝 없이 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 이상의 초대형 모델에 대한 검증은 향후 과제로 남아 있다.

실무 활용

기업용 LLM이나 특정 도메인 특화 모델을 개발할 때, 사전 학습된 모델을 바로 미세 조정하기보다 고품질 데이터를 활용한 미드 트레이닝 단계를 거치는 것이 성능 극대화에 유리하다.

수학/과학 등 전문 지식 추론이 필요한 도메인 특화 모델 개발
기존 베이스 모델의 코딩 능력을 저비용으로 강화하고자 하는 경우
강화 학습(RL) 적용 전 모델의 기초 체력을 다지기 위한 파이프라인 설계
장기 문맥 처리 능력을 유지하면서 특정 능력을 주입해야 하는 시나리오

코드 공개 여부: 공개

코드 저장소 보기

키워드

Mid-training(미드 트레이닝)Reinforcement Learning(강화 학습)Reasoning(추론)CKA(표현 유사도 분석)Weight Divergence(가중치 발산)

용어 해설

미드 트레이닝(Mid-Training): — 대규모 언어 모델 학습 과정에서 일반적인 Pre-training과 최종 Alignment 단계 사이에 수행되는 중간 학습 단계이다. 고품질의 도메인 특화 데이터를 주입하여 모델의 추론 능력을 비약적으로 향상시키고 후속 강화 학습을 위한 토대를 마련하는 중요한 역할을 한다.
중심 커널 정렬(CKA): — 서로 다른 신경망 레이어의 은닉 상태(Hidden State) 간의 유사성을 측정하는 지표이다. 모델의 내부 표현 공간이 학습 단계에 따라 어떻게 변화하거나 보존되는지 기하학적으로 분석하는 데 사용되며, 1에 가까울수록 표현 구조가 동일함을 의미한다.
그룹 상대 정책 최적화(GRPO): — 별도의 가치 모델(Critic) 없이 그룹 내 응답들의 상대적인 보상을 비교하여 정책을 업데이트하는 강화 학습 기법이다. 계산 효율성이 높으며 특히 수학이나 코드와 같이 정답이 명확한 추론 작업의 성능을 최적화하는 데 효과적이다.
사고 사슬(Chain-of-Thought): — 모델이 복잡한 문제에 대해 최종 정답을 내놓기 전 중간 추론 단계를 생성하도록 유도하는 기법이다. 미드 트레이닝을 통해 이러한 사고 패턴을 학습한 모델은 논리적 비약 없이 단계별로 문제를 해결하는 능력을 갖추게 된다.
선형 병합(Linear Merge): — 두 개 이상의 모델 가중치를 특정 비율로 산술 평균하여 하나의 모델로 합치는 기법이다. 이 논문에서는 미드 트레이닝 과정에서 손실된 베이스 모델의 장기 문맥 처리 능력을 복구하기 위해 베이스 모델과 미드 트레이닝 모델을 병합하는 용도로 사용된다.