[논문 리뷰] s1: Simple test-time scaling

Niklas Muennighoff, Zitong Yang외 8명|ArXiv.org|2025. 01. 31.

Fault Detection and Control Systems인용 수 10

한 줄 요약

본 논문은 테스트 시점 스케일링에 대한 최소한의 접근법을 제시한다: 1k 개의 추론 샘플에 대해 모델을 미세 조정하고 예산 강제(budget forcing)를 적용해 사고 시간을 제어하며, OpenAI o1-preview에 비해 경쟁력 있는 추론 성능과 데이터 효율성을 달성하고 오픈 소스 코드와 데이터를 제공한다.

ABSTRACT

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

연구 동기 및 목표

• 테스트 시점 컴퓨트가 언어 모델의 추론 성능을 향상시킬 수 있는 이유를 동기 부여한다.
• 다양하고 어려운 질문과 추적(trace)을 포함하는 작고 고품질의 추론 데이터셋(s1K)을 만든다.
• 생각 지속 시간을 제어하고 답변을 개선하기 위한 간단한 테스트 시점 개입(budget forcing)을 시연한다.
• 1K 샘플에서의 미세 조정이 강력하고 샘플 효율적인 추론 성능을 산출함을 보여준다.
• 재현과 추가 연구를 가능하게 하도록 오픈 소스 데이터, 모델 및 코드를 제공한다.

제안 방법

• 품질, 난이도, 다양성을 선정 기준으로 삼아 다양한 출처에서 59K 문제 풀을 큐레이션한다.
• 모델 기반 난이도 평가와 MSC 기반 도메인 다양성을 사용해 1K 개의 고품질, 다양하고 어려운 샘플(s1K)로 축소한다.
• s1K에서 Qwen2.5-32B-Instruct를 감독 학습(SFT)으로 미세 조정해 16대의 H100 GPU에서 26분 만에 s1-32B를 얻는다.
• 사고를 제어하기 위한 테스트 시간 중 예산 강제를 도입: (i) end-of-thinking 토큰을 첨가하여 사고를 종료, (ii) Wait를 첨가하여 더 많은 탐색을 장려.
• 순차적(budget forcing) 대 병렬(다수결) 방식의 테스트 시점 스케일링을 평가하고 벤치마크와 비교한다.
• 프로젝트 저장소에서 이용 가능한 오픈 소스 데이터, 가중치, 코드 등을 사용한다.

실험 결과

연구 질문

• RQ1 최소한의 데이터 효율적 접근법이 추론 과제에 대해 강력한 테스트 시점 스케일링을 달성할 수 있는가?
• RQ2 데이터 세트의 품질, 난이도, 다양성이 추론에 대한 지시 학습(instruction-finetuning)의 효과에 어떻게 영향을 미치는가?
• RQ3 순차적 테스트 시점 스케일링이 다수결과 같은 병렬 방법보다 더 효과적인가?
• RQ4 도전적인 추론 벤치마크에서 테스트 시점 컴퓨트가 증가함에 따라 budget forcing으로 인한 성능 향상은 어느 정도인가?
• RQ5 1K 샘플 s1K 데이터셋이 더 큰 풀과 비교하여 데이터 효율성과 경쟁력 있는 성능을 달성하는 데 어떤 차이가 있는가?

주요 결과

• 32B 모델의 1K 샘플 감독 학습은 MATH 및 AIME24에서 o1-preview에 경쟁력 있는 성능을 산출한다.
• 예산 강제는 제어 가능한 테스트 시점 컴퓨트를 가능하게 하고 모델 검증 및 더 긴 탐색을 촉진함으로써 추론을 향상시킨다.
• s1-32B는 사고 토큰이 더 허용될수록(수익 감소 직전까지) 성능이 증가하는 강한 테스트 시점 스케일링을 보인다.
• 이 접근법은 데이터 효율적이다: 1K 샘플에서의 학습이 많은 베이스라인보다 우수하고 더 큰 풀보다 훨씬 적은 예제로도 달성한다.
• 난이도, 다양성, 품질의 균형을 맞춘 데이터 선택이 중요하다; 무작위 또는 지나치게 긴 추적만으로는 제안된 세 가지 기준 선택에 비해 성능이 떨어진다.
• 1K 샘플로 학습된 s1-32B는 AIME24에서 Gemini 2.0 Thinking과 거의 비슷한 성능을 보이고 오픈 소스이다.