[논문 리뷰] rStar-Math: Small LLMs Can Master Math Reasoning with Self-Evolved Deep Thinking

Xinyu Guan, Li Lyna Zhang외 6명|arXiv (Cornell University)|2025. 01. 08.

Machine Learning and Data Classification인용 수 7

한 줄 요약

본 논문은 소형 언어 모델이 몬테 카를로 트리 탐색(MCTS)을 통한 자기 진화식 심층 사고, 코드-강화 CoT 데이터 합성 방법, 그리고 쌍별 순위를 통해 학습된 프로세스 선호 모델을 활용하여 OpenAI o1 스타일의 수학 추론에 도달하거나 이를 능가할 수 있음을 보여준다.

ABSTRACT

We present rStar-Math to demonstrate that small language models (SLMs) can rival or even surpass the math reasoning capability of OpenAI o1, without distillation from superior models. rStar-Math achieves this by exercising "deep thinking" through Monte Carlo Tree Search (MCTS), where a math policy SLM performs test-time search guided by an SLM-based process reward model. rStar-Math introduces three innovations to tackle the challenges in training the two SLMs: (1) a novel code-augmented CoT data sythesis method, which performs extensive MCTS rollouts to generate step-by-step verified reasoning trajectories used to train the policy SLM; (2) a novel process reward model training method that avoids naïve step-level score annotation, yielding a more effective process preference model (PPM); (3) a self-evolution recipe in which the policy SLM and PPM are built from scratch and iteratively evolved to improve reasoning capabilities. Through 4 rounds of self-evolution with millions of synthesized solutions for 747k math problems, rStar-Math boosts SLMs' math reasoning to state-of-the-art levels. On the MATH benchmark, it improves Qwen2.5-Math-7B from 58.8% to 90.0% and Phi3-mini-3.8B from 41.4% to 86.4%, surpassing o1-preview by +4.5% and +0.9%. On the USA Math Olympiad (AIME), rStar-Math solves an average of 53.3% (8/15) of problems, ranking among the top 20% the brightest high school math students. Code and data will be available at https://github.com/microsoft/rStar.

연구 동기 및 목표

• 소형 언어 모델이 더 큰 모델의 증류 없이도 최첨단 수학 추론 벤치마크에 맞먹거나 능가할 수 있음을 입증한다.
• 수학 추론을 위한 정책 모델과 보상 모델을 점진적으로 개선하는 자기 진화 워크플로우를 개발한다.
• 신뢰할 수한 단계별 궤적을 생성하기 위해 각 단계에서 파이썬 코드도 함께 생성하는 코드 증강 코-톳 데이터 합성 방법을 제안한다.
• 쌍별 순위를 통해 학습된 프로세스 선호 모델(PPM)을 도입하여 단계 수준의 밀집하고 신뢰할 수 있는 보상을 제공한다.
• 다양한 데이터셋과 모델 규모에서 경험적 이득을 보여주고, 수학 벤치마크에서 대형 프런티어 모델에 근접하거나 이를 능가한다(벤치마크별).

제안 방법

• 정책 SLM과 프로세스 보상 모델(PRM)을 사용하여 수학 추론에 대한 심층 사고를 수행하기 위해 몬테 카를로 트리 탐색(MCTS)을 사용한다.
• 각 단계에서 Python 코드도 생성되는 코드 증강 CoT 생성을 도입한다; 중간 단계의 유효성을 보장하기 위해 코드 실행이 가능한 생성만 보관한다.
• Q-값을 할당하기 위해 광범위한 MCTS 롤아웃으로 단계 품질을 주석하고, 종료 및 PRM-강화 전략을 사용하여 단계 점수를 정교하게 다듬는다.
• 노이즈가 많은 단계별 점수에 직접 의존하지 않고 고-저 Q 단계에 대해 쌍대 Bradley-Terry 순위로 학습된 프로세스 선호 모델(PPM)을 학습한다.
• 747k 수학 문제의 시드 데이터세트에서 시작하여 정책 SLM과 PPM을 점진적으로 강화하는 네 차례의 자기 진화 루프를 구현한다.
• 1.5B–7B SLM으로 MATH, AIME, AMC, Olympiad Bench 및 기타 벤치마크를 평가하고, OpenAI o1 및 다른 바깥 벤치마크와 비교한다.

실험 결과

연구 질문

• RQ1 소형 LLM이 더 큰 모델의 증류 없이도 프런티어 수학 추론에 도달할 수 있는가?
• RQ2 정책 모델과 보상 모델의 자기 진화가 다중 단계 수학 문제 해결의 격차를 어떻게 줄이는가?
• RQ3 코드-증강 CoT 데이터 합성이 추론 단계의 추적 품질을 개선하는가?
• RQ4 쌍대 순위로 학습된 프로세스 선호 모델이 수학 추론을 위한 신뢰할 수한 단계 수준 보상을 제공하는가?
• RQ5 다양한 수학 벤치마크에서 MCTS 궤적의 증가가 성능에 어떤 영향을 미치는가?

주요 결과

• rStar-Math는 7B 규모 모델을 최첨단 수준의 수학 벤치마크로 끌어올리며, 예를 들어 MATH 점수가 OpenAI o1에 상응하거나 그 이상으로 상승한다.
• MATH에서 Qwen2.5-Math-7B는 64 trajectories로 58.8%에서 90.0%로 향상되어 o1-preview를 능가하고 o1-mini와 일치한다.
• AIME 2024에서 rStar-Math는 평균 53.3%(8/15 문제)로 상위 20%의 가장 똑똑한 고등학생 수학 학생들 사이에 속한다.
• 수백만 개의 합성 해답으로 구성된 네 차례의 자기 진화는 정책 모델과 PPM을 점차 강화하여 747k 문제 중 90.25%로 커버리지를 확장한다.
• Python 실행 및 MCTS 기반 Q-값 주석이 있는 코드-증강 CoT는 중간 단계 오류를 줄이고 궤적 품질을 향상시킨다.
• 쌍대 순위로 학습된 프로세스 선호 모델(PPM)은 신뢰할 수 있는 단계 수준 가이던스를 제공하며 아블레이션에서 기저 보상 모델 접근법보다 우수하다.