자기생성 데이터로의 미드-트레이닝이 언어 모델의 Reinforcement Learning을 향상시키는 방법

RL의 효과는 바탕이 되는 데이터의 다양성에 좌우된다. 본 연구는 Polya의 문제해결 휴리스틱에 따라 자기생성 데이터의 다양한 해법을 중간 훈련에서 제시하고 이를 통해 RL 이후의 추론 능력을 넓히며, 수학적 추론뿐 아니라 코드 생성 및 서사적 추론에서의 일반화 이점을 보여준다.

출발점: RL은 base 모델의 priors에 의존한다. 단일 정답 해법만 제시하는 mid-training 데이터는 RL의 탐색 폭을 제한한다. 해결 원리: Polya의 휴리스틱에 따라 같은 문제에 대해 다양한 합리적 해법을 유도하고, 이를 바탕으로 각 해법의 경로를 포함한 다중 트레이닝 데이터를 구성한다. 이로써 mid-training에서 다중 모드(next-token 분포)가 형성되고, RL에서 이 모드들을 결합하는 방향으로 정책 업데이트가 유도된다. 달라지는 점: N-modal 분포 하에서의 업데이트는 하나의 모드에 과도하게 집중되는 것을 방지하고, RL이 여러 해법을 하나의 응답에 통합하는 능력을 촉진한다. 실험적으로는 해법의 다양성이 커질수록 pass@k의 상향이 지속적으로 나타나며, rollout 그룹 크기 g가 학습한 n에 맞을 때 성능이 더 높게 나타난다.

1. 전체 접근: 미드-트레이닝에서 각 질문 x에 대해 다수의 correct 솔루션 y(i)을 제공하는 n개의 해법 경로를 구성한다. 2) 데이터 구성: Dh(heuristic description)와 Eh(few-shot exemplars)로 구성된 각(h, x) 조합에 대해 base 모델 π를 조건부로 샘플링하고, 128개의 후보 응답을 생성한다. 3) 검증 및 선별: rule-based verifier(MATH-Verify 등)로 정답 여부를 판정하고, reward model Rϕ로 해리스를 따라가는지 점수화하여 최상위 응답을 최종 heuristic-specific 솔루션으로 선택한다. 4) MT 목표: L_MT(θ) = - E(x, Yx) [ (1/n) Σ_i Σ_t log πθ(y(i)t | y(i){<t}, x) ]. 이는 서로 다른 해결 전략에 대해 prefix에서 다수의 높은 확률 모드를 갖도록 하는 것을 목표로 한다. 5) RL 학습: DAPO-Math-17k 데이터셋과 GRPO 알고리즘을 사용해 RL을 수행하며, vLLM으로 롤아웃 생성과 평가를 가속한다. 6) 평가: six 수학 문제 벤치마크에서 pass@k를 측정하고, HumanEval/MuSR 같은 비수학 도메인으로 일반화 성과를 평가한다.

메인 벤치마크: pass@1은 n=64에서 평균 11.50%로 향상되었고, zero-shot 대비 평균 1.25% 포인트 증가를 관찰했다. pass@64의 평균은 46.30%(zero-shot)에서 STaR 46.32%, DPólya n=64에서 48.17%로 상승했다. Math-500에서 pass@64는 87.76%에서 88.94%로 증가했고, AIME 2025는 12.84%에서 18.66%로, AMC 2023은 83.49%에서 85.18%로, OlympiadBench은 42.13%에서 43.57%로 개선됐다. Out-of-domain: HumanEval에서 52.34%(n=32), 52.82%(n=64)로 Vanilla RL 대비 상승, MuSR에서도 각각 56.94%, 57.36%로 개선됐다. Distillation 대비: DPólya(n=16)에서 Vendi Score가 13.81로 측정, Distill+RL 대비 다변성에서 우위가 확인됐다. RL 결과는 rollout 그룹 크기 g=8에서 n=8이 세 가지 벤치마크에서 최상이며, g와 n의 매칭이 성능에 유의미한 영향을 준다.

아키텍처: base 모델로 Llama 3.2–3B–Instruct를 사용하고, 미드-트레이닝 데이터 DPólya를 구성한 뒤 RL에 적용한다. 학습은 4×NVIDIA H100에서 수행되며, 트레이닝은 한 에폭 444 스텝으로 고정하고, Effective batch size를 n에 따라 조정한다. MT 데이터 생성은 Dh, Eh를 포함하는 64개의 휴리스틱에 대해 각 x에 대해 128개의 샘플을 생성하고, Math-Verify와 보상 모델(Rϕ)을 통해 후보를 필터링한다. L_MT은 θ의 파라미터 업데이트의 음의 로지니-로그-가능도 평균으로 정의되며, 다중 해법 모드(y(i))마다의 로그가능도 합의 음수 평균을 손실로 사용한다. RL은 GRPO를 사용하고, DAPO-Math-17k를 롤아웃 데이터로 활용한다. 이론적 분석으로 Theorem 4.1은 단일 모드(unimodal)와 N-modal 분포에서의 정책 그래디언트 업데이트의 1차 확률 변화 차이를 도출하고, 다중 모드일수록 여러 접근법의 조합이 촉진됨을 보인다. 다중 해법을 통해), RL은 하나의 응답에 여러 해결 전략을 통합하는 경향이 증가한다. 벤치마크 외 일반화 실험은 HumanEval과 MuSR에서 수행되었고, Qwen 기반 모델에서도 유사한 경향을 확인했다.

A.1 한계: 휴리스틱은 수학 중심으로 구성되어 도메인 일반화에 대한 완전한 범주를 보장하지 않는다. BLayer 연구에서 제시된 분류 체계가 본 논문의 도메인에 완전히 일반화되지는 않는다. 미드-트레이닝 예산이 RL의 이득에 영향을 미치며, pass@1의 개선폭은 상대적으로 작다(주로 pass@64 수치에서 큰 효과). 학습 데이터의 자동 생성 특성상 인간 주석의 보조가 없는 경우 인간 수준의 일반화에 한계가 있을 수 있다.