s1: Simple test-time scaling 내용 정리

 

Introduction

최근 LLM 기반 응용들이 복잡한 추론 문제를 해결하기 위해 test-time scaling이라는 새로운 접근을 시도하고 있다. 이는 학습된 모델을 그대로 두면서, 테스트 시점에서 추가 연산을 허용하여 성능을 더 높이는 개념이다. 예컨대, OpenAI가 선보인 o1 모델은 테스트 시점에 더 많은 연산을 사용해 성능을 꾸준히 향상시켜, 다양한 수학·과학 문제에서 좋은 결과를 보였다. 그러나 해당 기법이 비공개였기에 많은 연구 그룹이 유사한 시도를 했으나, 흔히 RL이나 상당한 양의 data를 사용하는 방식이었다.

이번 글에서는 Simple Test-time Scaling을 표방하는 새로운 방법론과, 이를 이용해 만든 오픈 모델인 s1-32B를 다룬다. 핵심 포인트는 다음과 같다.

• 1,000개의 고품질, 고난도, 다분야 reasoning data(s1K)를 만들었다.
• 이 작은 데이터로 SFT(Supervised Finetuning)만 진행하고, budget forcing이라는 단순한 기법으로 test-time scaling을 구현했다.
• 결과적으로 기존 모델 Qwen2.5-32B-Instruct 대비 적은 자원으로 의미 있는 reasoning 성능 향상과 test-time scaling까지 확인했다.

---

s1K: 1,000개 reasoning 데이터셋 구축

전체 데이터 수집

우선 59K 규모의 문제 집합을 모았다. 주요 출처는 다음과 같다.

1. NuminaMATH 등 대규모 수학 문제.
2. AIME 등 실제 대회 문제.
3. 올림피아드 대회, AGIEval, OlympicArena 등 어려운 과목 범위.
4. 별도의 확장 데이터로 s1-prob(박사과정 확률 문제), s1-teasers(난이도 높은 퍼즐) 등 소량 추가.

모두 Google Gemini를 통해 reasoning trace와 해답을 생성했다.

1K 최종 선별

처음 수집한 59K 중에서, (1) 품질이 낮거나 오류가 있는 예시 제거, (2) Qwen2.5-7B, 32B 모델로 풀어봤을 때 너무 쉽게 풀리는 문제 제외, (3) Math Subject Classification 기반으로 다양한 분야가 골고루 들어가도록 균형 있게 샘플링하여 최종 1,000개의 s1K를 완성했다. 이때 난이도, 다양성, 품질이라는 세 가지 기준을 종합적으로 고려했다.

---

Test-time Scaling과 Budget Forcing

 

기존 아이디어

보통 모델의 성능을 높이는 전략은 크게 parallel(다중 시도 뒤 투표)와 sequential(단계적으로 더 길게 추론) 방식으로 나뉜다. parallel 방식인 majority voting이나 Best-of-N은 샘플을 여러 번 생성한 뒤, 뭔가 기준(투표, reward 등)으로 최종 해답을 고른다. 반면 sequential은 한 번에 긴 reasoning trace를 생성하거나, 전 단계 생성 결과를 토대로 추가 추론을 이어나간다.

Budget Forcing

Budget forcing의 예시 Wait를 활용하여 resoning을 지속시킨다.

 

본 연구에서는 budget forcing이라는 간단한 decode-time 기법을 제안했다. 구체적으로:

1. 제한 토큰 수 이상으로 생각할 때: end-of-thinking 토큰을 강제로 삽입해 reasoning 단계를 조기에 종료하고 답을 내도록 한다.
2. 제한 토큰 수보다 더 많이 생각하게 하고 싶을 때: end-of-thinking 토큰 생성을 막고, “Wait” 같은 단어를 reasoning 중간에 추가해 모델이 생각을 이어가게 한다.

이 방식은 실제로 test-time에서만 적용되므로 모델 파라미터나 학습 데이터는 동일하지만, 달리 주어진 “시간(토큰)”만큼 더(또는 덜) 생각하도록 컨트롤이 가능하다. 실험적으로 이 기법을 통해 AIME24, MATH500 등 competition math에서 reasoning 품질이 뚜렷이 증가하는 것을 확인했다.

---

s1-32B: 결과 및 비교

모델 개요

• 베이스 모델: Qwen2.5-32B-Instruct
• 데이터: s1K(1,000개의 복잡 추론 예시)
• 학습: 기본 SFT만 사용, 약 26분간(16 x H100) 수행
• 테스트 시점: budget forcing을 통해 reasoning 길이(토큰 수)에 따른 성능 변화를 확인

성능

• AIME24에서 최대 57%, MATH500에서 최대 93% accuracy를 달성했다.
• GPQA Diamond에서도 59~60%대 성능을 보여, 다양한 분야 과학문제에 대해서도 괜찮은 추론 성능을 확인했다.
• 특히 s1-32B는 오픈AI o1-preview와 비슷하거나 더 나은 수치를 보였다.
• DeepSeek 등 대규모 RL 기반 모델 대비 데이터 효율이 훨씬 높았다.

1,000개의 carefully curated 데이터로도 충분히 견고한 reasoning 모델을 만들 수 있고, test-time scaling을 단순 decode-level 개입(budget forcing)으로 구현해, 적은 자원으로도 o1 같은 비공개 고성능 모델에 준하는 성능을 시현했다.

---

결론

본 연구는 간결함이 실제로도 강력한 방법이라는 점을 보여준다. 1,000개의 소규모 but 고품질 reasoning 데이터와, “Wait”만 붙여 더 생각하도록 유도하는 단순 방법인 budget forcing으로 test-time scaling을 성공적으로 재현했다. 추가적으로 다음과 같은 향후 연구가 있을 수 있다.

• 대형 모델과의 결합: 현재는 32B급 모델에 시현했지만, 더 큰 모델에서도 적은 finetuning만으로 큰 효과를 낼지 확인이 필요하다.
• 다른 test-time scaling 방안: parallel이나 tree-based search 같은 고급 기법과 budget forcing을 섞으면 더 큰 시너지를 낼 수 있다.
• 다양한 도메인 적용: 경쟁 수학 문제에 국한되지 않고, 의료/법/연구 논문 등 복잡한 도메인에 적용 가능성을 탐색해야 한다.