s1: Simple test-time scaling

Test-time scaling is a promising new approach to language modeling that uses extra test-time compute to improve performance. Recently, OpenAI's o1 model showed this capability but did not publicly share its methodology, leading to many replication efforts. We seek the simplest approach to achieve test-time scaling and strong reasoning performance. First, we curate a small dataset s1K of 1,000 questions paired with reasoning traces relying on three criteria we validate through ablations: difficulty, diversity, and quality. Second, we develop budget forcing to control test-time compute by forcefully terminating the model's thinking process or lengthening it by appending"Wait"multiple times to the model's generation when it tries to end. This can lead the model to double-check its answer, often fixing incorrect reasoning steps. After supervised finetuning the Qwen2.5-32B-Instruct language model on s1K and equipping it with budget forcing, our model s1-32B exceeds o1-preview on competition math questions by up to 27% (MATH and AIME24). Further, scaling s1-32B with budget forcing allows extrapolating beyond its performance without test-time intervention: from 50% to 57% on AIME24. Our model, data, and code are open-source at https://github.com/simplescaling/s1

한 줄 요약

s1-32B 모델은 1,000개의 s1K 데이터셋과 budget forcing 기법을 통해 o1-preview를 27% 초과하는 수학 문제 해결 성능을 달성했다.

핵심 기여도

s1K 데이터셋: 1,000개의 문제-추론 트레이스 쌍으로, difficulty, diversity, quality 기준을 통해 선별.
budget forcing: test-time compute를 제어하는 간단한 기법으로, "Wait" 삽입 및 강제 종료를 통해 추론 과정 조절.
s1-32B 모델: Qwen2.5-32B-Instruct를 s1K로 fine-tuning한 모델로, AIME24에서 o1-preview 대비 27% 성능 향상.
open-source 공개: 모델, 데이터, 코드를 https://github.com/simplescaling/s1 에 공개.

핵심 아이디어

기존 연구는 test-time scaling을 위해 복잡한 강화학습(RL)이나 대규모 데이터를 사용했으나, 본 연구는 가장 단순한 접근법을 추구했다. 핵심 아이디어는 **1,000개의 소규모 but 고질적인 데이터셋(s1K)**과 **test-time compute를 조절하는 budget forcing 기법**을 결합하는 것이다. s1K는 difficulty, diversity, quality 세 가지 기준을 통해 엄선된 문제-추론 트레이스 쌍으로 구성되어 있으며, 이는 고효율 학습을 가능하게 한다. budget forcing은 모델이 추론을 종료하려 할 때 "Wait"를 반복 삽입하거나, 특정 토큰 수를 넘기면 강제 종료함으로써, **추론 과정을 조절하고 오류를 수정**하도록 유도한다. 이는 복잡한 알고리즘 없이도 test-time scaling을 달성할 수 있음을 보여준다.

기술적 접근법

**데이터셋**: s1K는 1,000개의 문제-추론 트레이스 쌍으로 구성. Gemini Thinking Experimental에서 추출.
**모델**: Qwen2.5-32B-Instruct를 기반으로, s1K 데이터셋을 사용한 supervised fine-tuning(SFT) 수행.
**훈련 환경**: 16개의 H100 GPU를 사용해 26분 동안 훈련.
**test-time 기법**: budget forcing.
(I) 토큰 수가 제한을 넘으면 강제 종료.
(II) "Wait" 삽입으로 추론 시간 연장.
**성능 측정**: AIME24, MATH500, GPQA Diamond 데이터셋 사용.

주요 결과

**AIME24**: s1-32B는 o1-preview 대비 27% 성능 향상.
**AIME24에서 test-time scaling**: 50% → 57% 성능 향상 (test-time interventional 없이).
**MATH500, GPQA Diamond**: OpenAI 기준과 동일한 샘플 사용.
**데이터셋 실험**: 59K 샘플 풀에서 1K 샘플만 선택해 훈련했음에도, 랜덤 선택 또는 최대 다양성 선택보다 30% 이상 성능 우수.

의의 및 한계

s1-32B는 test-time scaling을 달성하기 위해 **복잡한 RL이나 대규모 데이터를 사용하지 않고도 뛰어난 성능**을 보인 점에서 의의가 있다. 특히, 1,000개의 소규모 데이터셋으로도 강력한 추론 모델을 구축할 수 있음을 보여주며, **샘플 효율성** 측면에서 기존 연구를 뛰어넘는다. 또한, budget forcing은 test-time compute를 정확히 제어할 수 있는 **단순하면서도 효과적인 기법**으로, 다양한 모델에 적용 가능하다는 장점이 있다. 한편, 본 연구는 수학 문제에 초점을 맞춘 관계로, 다른 분야(예: 과학, 언어)로의 확장 가능성은 명시되지 않았으며, 추론 과정의 질적 분석도 부족한 점이 한계로 지적된다.

실용적 활용

s1-32B와 budget forcing 기법은 **경쟁 수학 문제 해결, 자동 평가 시스템, 교육 AI** 등에 활용될 수 있다. 특히, 복잡한 알고리즘 없이도 test-time 성능을 향상시킬 수 있으므로, **자원이 제한된 환경**에서도 유용하게 사용될 수 있다. 또한, open-source 공개를 통해 연구자들이 다양한 분야에 적용해볼 수 있는 기회를 제공한다.