Mid-Training with Self-Generated Data Improves Reinforcement Learning in Language Models

기술 요약: 자체 생성 데이터를 활용한 중간 훈련이 언어 모델의 강화 학습을 개선함

문제 제기
대규모 언어 모델 (LLM) 의 최근 발전은 추론 시간 확장 (inference-time scaling) 을 통해 추론 능력을 향상시키는 강력한 기법으로서 강화 학습 (RL) 을 부각시켰습니다. 그러나 RL 의 효과는 RL 훈련 이전에 모델에 확립된 사전 지식 (priors) 에 크게 의존합니다. 기존 문헌은 검증 가능한 보상을 통한 강화 학습 (RLVR) 이 종종 '분포 sharpening'을 수행하여, 기저 모델에 이미 존재하는 행동을 선택적으로 증폭시킬 뿐 진정한 새로운 능력을 유도하지는 않는다고 시사합니다. 모델이 단일 작업에 대한 다양한 문제 해결 접근법에 노출되지 않은 경우, 후속 RL 훈련은 성능 향상에 어려움을 겪거나 특정 평가 지표 (예: pass@k) 하에서는 오히려 성능이 저하될 수 있습니다. 본 연구가 해결하는 핵심 과제는 RL 적용 전에 기저 모델에 다양한 추론 전략의 풍부한 세트를 효과적으로 갖추게 하여, RL 과정이 단순히 기존 행동의 좁은 집합을 강화하는 것이 아니라 이러한 전략들을 구성하고 정제할 수 있도록 하는 것입니다.

방법론
저자들은 사전 훈련과 RL 사이의 중간 단계로 작용하는 중간 훈련 (mid-training) 프레임워크를 제안합니다. 이 단계는 동일한 문제에 대한 여러 유효한 해결 궤적을 모델에 노출시키기 위해 다양한 자체 생성 데이터를 활용합니다.

1. 휴리스틱 유도 데이터 생성:

   • 이 프레임워크는 George Pólya 의 문제 해결 휴리스틱 (예: 유추, 분해 및 재구성, 역산) 에서 영감을 얻습니다.
   • 시드 데이터셋 (GSM8K) 의 각 질문에 대해, 기저 모델은 특정 휴리스틱 설명과 퓨샷 예시 (few-shot exemplars) 를 프롬프트로 받아 여러 후보 응답을 생성합니다.
   • 이 과정은 부트스트랩 방식으로 수행됩니다: 모델은 더 강력한 교사 모델로부터의 증류 (distillation) 에 의존하지 않고 자체 훈련 데이터를 생성합니다.
   • 필터링 및 선택: 생성된 응답은 규칙 기반 검증기 (예: Math-Verify) 를 사용하여 정확도로 필터링됩니다. 남은 후보들은 특정 휴리스틱 준수 여부에 따라 보상 모델 ($R_\phi$) 에 의해 점수가 매겨집니다. 각 질문 - 휴리스틱 쌍에 대해 가장 높은 점수를 받은 응답이 선택됩니다.
   • 결과 데이터셋 ($D_{Pólya}$) 은 질문당 $n$개의 서로 다른 해결 변형을 포함하며, 여기서 $n$은 적용된 휴리스틱의 수를 나타냅니다.

2. 중간 훈련 (지도 미세 조정):

   • 기저 모델은 $D_{Pólya}$에 대해 미세 조정됩니다. 질문과 단일 응답을 짝짓는 표준 SFT 와 달리, 이 접근법은 각 질문에 $n$개의 다양한 궤적을 짝짓습니다.
   • 목표는 이러한 여러 궤적에 걸쳐 음의 로그 가능도를 최소화하여, 모델이 분기점에서 다양한 다음 토큰 옵션에 무시할 수 없는 확률을 할당하도록 유도하는 것이며, 이는 효과적으로 N-모달 다음 토큰 분포를 생성합니다.

3. 강화 학습:

   • 중간 훈련 후, 모델은 별도의 데이터셋 (DAPO-Math-17k) 에서 그룹 상대 정책 최적화 (GRPO) 를 사용하여 RL 훈련을 받습니다.
   • 본 연구는 다양한 수의 휴리스틱으로 초기화된 모델 ($n \in \{1, 2, \dots, 64\}$) 을 제로샷, 바닐라 RL(기저 모델에 직접 적용), STaR+RL(자기 학습 추론가 후 RL) 과 같은 베이스라인과 비교합니다.

이론적 관점
본 논문은 중간 훈련이 정책 경사 업데이트에 미치는 영향을 이론적으로 분석합니다:

• 단일 모달 vs N-모달 체제: 단일 토큰에 대한 높은 확신을 가진 단일 모달 체제에서 정책 경사 업데이트는 미미한 변화 ($\propto \epsilon^2$) 를 초래합니다. 반면, $N$개의 지배적 모드 간에 확률 질량이 분할된 N-모달 체제에서는 업데이트가 보수적이지만 appreciable 한 변화 ($\propto 1/N$) 를 유도합니다.
• 전략의 구성: 저자들은 음의 이점 업데이트 하에서, 샘플링된 토큰에서 제거된 확률 질량이 주로 나머지 $N-1$개의 지배적 모드로 재분배됨을 증명합니다. 이 메커니즘은 모델이 RL 동안 단일 전략으로 수렴하는 대신 단일 응답 내에서 서로 다른 문제 해결 접근법을 결합하도록 학습하도록 장려합니다.

주요 결과
실험은 Llama 3.2–3B–Instruct 와 Qwen2.5–7B–Instruct 를 대상으로 Math-500, AIME 2024/2025, AMC 2023, HMMT 2025, OlympiadBench 등 여섯 가지 수학 추론 벤치마크 및 도메인 외 작업 (HumanEval, MuSR) 에서 수행되었습니다.

• 성능 향상: 다양한 자체 생성 데이터를 통한 중간 훈련은 RL 성능을 일관되게 향상시키며, 특히 높은 $k$값 (pass@64) 에서 두드러집니다.
  • Math-500 에서 pass@64 는 제로샷 (87.76%) 에서 $n=64$ (88.94%) 로 향상되었습니다.
  • AIME 2025 에서 pass@64 는 12.84% 에서 18.66% 로 증가했습니다.
  • 더 높은 추론 다양성 ($n=64$) 으로 중간 훈련된 모델은 모든 벤치마크에서 바닐라 RL 과 STaR+RL 베이스라인을 모두 능가했습니다.
• 다양성 대 양: 통제 실험은 고정된 훈련 예산 하에서 더 많은 문제에 대한 단일 접근법 학습보다 질문당 여러 접근법 학습 (더 적은 문제, 더 많은 접근법) 이 RL 에 더 큰 이점을 제공함을 보여주었습니다.
• 정확성은 필수적입니다: 잘못된 답으로 이어지는 다양한 접근법으로의 중간 훈련은 RL 성능을 향상시키지 못했으며, 이는 정확성이 다양성의 이점을 위한 전제 조건임을 나타냅니다.
• 증류보다 우월함: 자체 생성 데이터는 더 강력한 교사 모델 (QwQ-32B) 로부터 증류된 데이터 (Vendi Score 10.95) 에 비해 더 높은 다양성 (Vendi Score 13.81) 을 보였으며, 이로 인해 더 나은 다운스트림 성능이 달성되었습니다.
• 일반화: Pólya 의 휴리스틱이 수학 중심임에도 불구하고, 학습된 전략은 코드 생성 (HumanEval) 과 내러티브 추론 (MuSR) 으로 일반화되었으며, 다단계 추론 작업에서 상당한 향상을 보였습니다.

의의 및 주장
본 논문은 중간 훈련 동안 자체 생성 데이터를 통해 여러 문제 해결 접근법에 노출되는 것이 모델의 사전 지식을 근본적으로 형성하여, 후속 RL 이 이러한 전략들을 효과적으로 구성할 수 있게 한다는 것이 주요 기여라고 주장합니다.

• 개선 메커니즘: 본 연구는 RL 이 반드시 처음부터 '창발적'인 추론 능력을 생성하는 것이 아니라 기존 행동을 구성한다는 점을 시사합니다. 중간 훈련은 이러한 행동들이 다양하며 구성을 위해 사용 가능하도록 보장합니다.
• 실용적 유용성: 이 접근법은 인간 주석이나 더 강력한 교사 모델 없이 RL 을 개선하는 간단하고 효과적이며 확장 가능한 방법을 제공합니다.
• 이론적 통찰: 이 작업은 RL 의 경험적 성공과 이론적 이해 사이의 간극을 메우며, N-모달 분포가 모드 붕괴를 방지하고 복합 추론 체인의 학습을 촉진하는 방식을 보여줍니다.

저자들은 단일 샘플 정확도 (pass@1) 에 대해서는 겸손한 입장을 유지하며, 그들의 초점은 다양한 노출이 RL 행동에 어떻게 영향을 미치는지 이해하고 여러 샘플링 시도 (pass@k) 가 필요한 시나리오에서 성능을 향상시키는 데 있음을 지적합니다. 그들은 휴리스틱이 수학 중심이지만, 그 결과로 얻어진 다른 도메인에서의 일반화는 이 절차가 기저 모델 내의 잠재적 능력을 표면화하고 통합함을 시사한다고 인정합니다.