rStar-Math: Small LLMs Can Master Math Reasoning with Self-Evolved Deep Thinking

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\"ive 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.

한 줄 요약

rStar-Math는 MCTS와 자체 진화를 통해 소형 언어 모델의 수학 추론 성능을 OpenAI o1 수준으로 끌어올리는 시스템 2 스타일 접근법이다.

핵심 기여도

코드 증강된 사고 과정(CoT) 데이터 합성 방법을 도입하여 정확한 단계별 추론 트레이젝토리를 생성.
단계별 점수 주석 없이 더 효과적인 프로세스 선호 모델(PPM)을 학습하는 새로운 방법 제시.
정책 SLM과 PPM을 4라운드에 걸쳐 자체 진화시켜 수학 추론 능력을 향상.
Qwen2.5-Math-7B에서 58.8% → 90.0%, Phi3-mini-3.8B에서 41.4% → 86.4%로 성능 향상.

핵심 아이디어

rStar-Math는 기존의 대형 언어 모델(LLM)을 사용한 데이터 증강에 의존하지 않고, 소형 언어 모델(SLM)을 기반으로 자체 진화를 통해 수학 추론 능력을 향상시키는 새로운 접근법이다. 핵심 아이디어는 Monte Carlo Tree Search(MCTS)를 통해 "깊은 사고"를 유도하는 동시에, 정확한 단계별 추론을 생성하고 이를 기반으로 정책 모델과 PPM을 반복적으로 훈련하는 것이다. 이는 기존의 단계별 점수 주석이 필요 없는 PPM 학습과, 코드 실행을 통해 검증된 추론 트레이젝토리를 생성하는 CoT 데이터 합성 기법을 통해 가능하다. 특히, PPM은 Q-value 기반의 선호 쌍을 생성하여 정확도 높은 단계별 보상 예측을 가능하게 하며, 이는 기존의 ORM 및 PQM보다 우수한 성능을 보인다.

기술적 접근법

**MCTS 기반 추론**: 정책 SLM이 MCTS를 통해 단계별 추론을 생성하며, Python 코드 실행을 통해 검증.
**코드 증강된 CoT 데이터 합성**: MCTS rollout을 통해 생성된 단계별 추론 트레이젝토리 중, 성공적으로 실행된 코드만을 포함.
**PPM 학습**: Q-value 기반의 선호 쌍 생성과 pairwise ranking loss를 사용하여 단계별 보상 예측.
**자체 진화 레시피**: 4라운드에 걸쳐 정책 SLM과 PPM을 반복적으로 훈련하며, 747k 수학 문제에 대한 수백만 개의 합성 솔루션을 생성.
**하이퍼파라미터**: 8개의 추론 경로에서 Qwen2.5-Math-7B는 58.8% → 89.4%, 64개 경로에서는 90.0%로 성능 향상.

주요 결과

**MATH 벤치마크**: Qwen2.5-Math-7B에서 58.8% → 90.0% (o1-preview 대비 +4.5%), Phi3-mini-3.8B에서 41.4% → 86.4% (o1-preview 대비 +0.9%).
**AIME 2024**: 평균 53.3% (8/15) 문제 해결, o1-preview 대비 +8.7% 개선.
**자체 진화 효과**: 4라운드 후 정확도가 지속적으로 향상되며, 2라운드부터 GPT-4o를 초과.
**PPM 효과**: ORM 대비 PPM이 더 높은 정확도를 보이며, 특히 MATH와 올림피아드 벤치마크에서 우수한 성능.

의의 및 한계

rStar-Math는 소형 언어 모델이 대형 모델의 수학 추론 능력을 따라잡거나 초과할 수 있음을 입증하며, 기존의 데이터 증강에 의존하지 않는 새로운 학습 패러다임을 제시한다. 특히, 자체 생성된 데이터를 통해 정책 모델과 PPM을 반복적으로 훈련하는 방식은 수학 추론 분야에서 중요한 기술적 진보를 의미한다. 그러나, 현재는 747k 문제에만 적용되었으며, 더 어려운 문제나 다른 도메인으로 확장 가능성은 아직 검증되지 않았다. 또한, PPM의 Q-value 기반 학습이 여전히 일부 노이즈를 포함할 수 있다는 한계도 존재한다.

실용적 활용

rStar-Math는 수학 교육, 온라인 학습 플랫폼, 자동 문제 풀이 시스템 등에서 활용될 수 있다. 특히, 소형 모델을 사용해 고성능 수학 추론을 구현할 수 있으므로, 컴퓨팅 자원이 제한된 환경에서도 유용하게 사용될 수 있다. 또한, 자체 진화 방식은 다른 추론 도메인으로 확장 가능하며, 연구자들이 새로운 학습 전략을 개발하는 데 기초 자료로 활용될 수 있다.