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

 

arXiv 2025. [Paper] [Github]
Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, Tatsunori Hashimoto
Stanford University | University of Washington | Allen Institute for AI | Contextual AI
31 Jan 2025

개요

해당 논문은 Test-time scaling이라는 개념을 활용하여 언어 모델의 성능 향상시키는 방법을 탐구한다. 최근 OpenAI의 o1모델이 이 기술을 사용하여 뛰어난 성능을 보였으나, 구체적인 방법이 공개되지 않았다. 저자들은 가장 단순한 접근법을 통해 test-time scaling을 달성하는 방법을 찾고자 했으며, 이를 위해 Budget Forcing이라는 새로운 기법을 제안했다.

 

기여점

1. Test-time scaling의 단순한 접근법을 제안
   • 사전 학습된 모델을 대상으로 소규모 데이터(1,000개의 문제)로 지도 학습(Supervised Fine-Tuning, SFT)을 수행
   • 이후 Budget Forcing을 활용해 모델이 더 오랜 시간 "생각"하도록 유도하여 성능 향상
2. s1K 데이터셋 구축
   • 기존 수집된 59,029개 질문에서 품질(quality), 난이도(difficulty), 다양성(diversity)을 기준으로 1,000개의 문제를 선별
   • Google Gemini의 reasoning trace를 활용하여 데이터 정제
3. Budget Forcing 기법 제안
   • 모델의 reasoning token 길이를 조절하여 반드시 더 오래 생각하도록 유도하거나 조기 종료시킴
   • "Wait" 토큰을 추가하면 모델이 스스로 답을 검토하고 수정하는 경우가 많음
4. 경쟁 모델 대비 성능 향상 및 Sample Efficiency 입증
   • OpenAI o1-preview 대비 AIME24(수학 경시대회 문제)에서 최대 27% 성능 향상
   • Test-time scaling을 통해 학습 데이터 없이도 성능을 추가적으로 향상시킬 수 있음

Test-Time Scaling이란?

기존 언어 모델의 성능 개선 방법은 훈련(Train-Time)에서의 Scaling을 통해 모델 크기와 데이터 양을 증가시키는 것이었음 하지만, Test-time scaling은 추론(Inference) 과정에서 연산량을 증가시켜 성능을 향상하는 방법을 의미함
즉, 기존 모델을 더 효율적으로 활용할 수 있도록 추론 단계에서 더 깊은 reasoning을 유도하는 방법을 연구하는 것

데이터셋 구축: s1K

1단계

데이터 수집 기준

59K 데이터를 수집할 때, 세 가지 주요 원칙을 따름

데이터 필터링 3단계

1. 품질 (Quality):
   • 데이터셋이 고품질이어야 하며, 사람이 직접 샘플을 검토하여 문제 형식이 깨지거나 오류가 있는 데이터를 제거함
2. 난이도 (Difficulty):
   • 간단한 문제보다는 심층적인 reasoning이 필요한 문제를 포함해야 함
   • 특정 모델(Qwen2.5-7B, Qwen2.5-32B)로 문제를 풀어본 뒤, 쉽게 풀리는 문제는 제외
3. 다양성 (Diversity):
   • 여러 분야(수학, 물리, 통계, 생물학 등)를 포함하여 다양한 reasoning 방식이 필요하도록 구성

데이터 소스

6개의 다양한 데이터셋에서 문제를 수집하여 총 59,029개 질문을 확보
이 중 주요한 데이터셋은 다음과 같음

1) 수학 관련 데이터셋 (Mathematical Problem Solving)

• NuminaMATH (30,660개 샘플)
  • 온라인 수학 문제 사이트에서 수집한 고난이도 문제들
• AIME (미국 수학 경시대회, 1983~2021년 기출 문제 포함)
  • 미국 중·고등학교 수학 올림피아드 예선 문제
• OlympicArena (4,250개 샘플)
  • 다양한 올림피아드 수준의 문제 포함 (천문학, 생물학, 화학, 컴퓨터 과학, 지리학, 수학, 물리학 등)
• OmniMath (4,238개 샘플)
  • 경쟁적 수학 문제로 구성된 데이터셋

2) 표준화 시험 문제 데이터셋 (Standardized Tests)

• AGIEval (2,385개 샘플)
  • SAT, LSAT, GRE 등의 표준 시험 문제 포함
  • 영어, 법률, 논리 등 다양한 도메인을 포함

3) 새로운 데이터셋 구축 (Custom Dataset Creation)

• s1-prob (182개 샘플)
  • 스탠퍼드 대학교 통계학 박사 자격시험(PhD Qualifying Exams) 문제를 기반으로 수집
  • 확률 및 증명 문제 포함
• s1-teasers (23개 샘플)
  • 금융·트레이딩 면접에서 사용되는 어려운 두뇌 퍼즐(brain teaser) 문제 포함

저자들은 기존 14개의 데이터셋들에서 총 58,824개의 질문을 수집하였으며, 기존 데이터셋을 보완하기 위해 두 개의 새로운 데이터셋인 s1-prob와 s1-teasers를 만들었다. 전체 데이터 분포는 아래 표와 같다.

• Google Gemini Flash Thinking API를 이용해 reasoning trace를 생성하여 문제와 정답을 함께 구성
• 중복된 문제 제거 및 데이터 정제
• MATH500, GPQA Diamond, AIME24 등의 평가 데이터셋과 중복되지 않도록 8-gram 중복 제거 기법 적용

2단계

59K의 원본 데이터 중 최적의 1,000개 샘플을 선별하는 과정
이를 위해 3단계 필터링을 수행함

1. 품질 필터링 (Quality Filtering)

• API 호출 중 오류가 발생한 데이터 삭제 → 54,116개 샘플로 감소
• 포맷이 깨지거나 잘못된 수식, ASCII 아트, 깨진 이미지가 포함된 문제 제거 → 51,581개 샘플로 감소
• 정확한 reasoning trace가 있는 데이터만 유지

2. 난이도 필터링 (Difficulty Filtering)

두 가지 방법으로 난이도를 측정하여 어려운 문제만 선별

1. 모델 성능 기반 필터링
   • Qwen2.5-7B 및 Qwen2.5-32B 모델을 사용하여 문제를 풀어봄
   • 둘 중 하나라도 정답을 맞힐 경우 해당 문제 제거 → 쉬운 문제 배제
   • → 24,496개 샘플로 감소
2. Reasoning Trace 길이 기반 필터링
   • 문제를 해결하는 reasoning trace의 토큰 수(token length)가 길수록 난이도가 높다고 판단
   • 짧은 reasoning을 가진 문제를 추가적으로 제거

💡 왜 두 가지 모델을 사용했을까?

• 단순히 한 모델만 사용하면 우연히 오답을 낸 쉬운 문제가 남을 가능성이 있음
• 두 모델 모두 맞힌 문제만 배제함으로써 더 정확한 난이도 조절 가능

3. 다양성 필터링 (Diversity Filtering)

다양한 도메인을 포함하도록 문제를 선별

• 수학 분야뿐만 아니라 물리, 생물, 경제 등 다양한 분야에서 문제 선택
• Mathematics Subject Classification (MSC) 시스템을 활용하여 도메인 태깅
  • 예: 기하학(geometry), 확률(probability), 양자역학(quantum mechanics), 경제학(economics) 등
• 최종적으로 50개 이상의 분야에서 균형 있게 문제 선택
• 최종적으로 1,000개의 문제를 확보

  ---

Test-time scaling - Method

 

저자들은 test-time scaling 방법을 두 가지로 분류하였다.

1. Sequential scaling: 나중의 계산이 이전 계산에 의존하는 경우 (ex. 긴 추론 과정)
2. Parallel scaling: 계산이 독립적으로 실행되는 경우 (ex. 다수결 투표)

저자들은 sequential scaling에 집중하였는데, 나중의 계산이 중간 결과를 기반으로 구축될 수 있기 때문에 더 깊은 추론과 반복적 정제가 가능하고, 따라서 더 나은 scaling이 가능하기 때문이다. 저자들은 새로운 sequential scaling 방법과 이를 벤치마킹하는 방법을 제안하였다.

---

Budget Forcing: Test-time Scaling 기법

Budget Forcing(BF)은 논문에서 제안한 Test-time Scaling 기법으로, 모델이 특정 문제를 푸는 데 사용하는 연산량을 조절하여 reasoning 성능을 향상시키는 방법

핵심 아이디어

• 모델이 문제를 해결할 때 충분한 reasoning을 수행할 수 있도록 제어하는 기법
• "얼마나 오래 생각할지(Thinking Time)"를 조절하여 모델의 reasoning 능력을 향상시킴
• 추론 도중 reasoning이 부족하거나 실수가 발생할 가능성이 있는 경우 추가적인 reasoning을 강제하여 정답률을 높임

 

---

왜 Budget Forcing이 필요한가?

일반적인 언어 모델의 한계
모델이 reasoning 없이 빠르게 답을 내려고 하는 경향이 있음.

• 특히, 모델이 빠르게 종료할 수 있는 경우, 충분히 검토하지 않고 "비효율적인 답변"을 생성할 수 있음.

Test-time Scaling이 필요함

• 언어 모델이 reasoning을 깊게 수행할수록 성능이 좋아지는 경향이 있음.
• 하지만, 모델이 reasoning을 멈추는 시점을 적절히 조절하지 않으면, 충분한 reasoning을 수행하지 못할 수 있음
• Budget Forcing을 사용하면 모델이 더 오랜 시간 reasoning을 수행하도록 유도할 수 있음

---

Budget Forcing 동작 방식

Budget Forcing은 두 가지 방식으로 모델이 reasoning을 더 많이 하거나 더 적게 하도록 조정할 수 있음

Maximum Budget: 제한된 thinking time 내에서 reasoning 수행

• 제한된 reasoning token 개수를 넘지 않도록 강제 종료
• 특정 토큰 수 이상을 사용하지 않도록 설정하면, 모델이 빠르게 reasoning을 마치고 정답을 도출하게 됨
• 활용 예시:
  • 모델이 너무 길게 reasoning을 하면 불필요한 내용이 포함될 수 있음
  • 이런 경우 "최대 2,000개의 reasoning token을 사용하도록 제한"하는 방식으로 적용 가능

Minimum Budget: 추가적인 reasoning 강제

• 모델이 스스로 reasoning을 멈추려고 할 때 강제로 reasoning을 더 수행하도록 유도
• 구체적인 방법:
  • 모델이 reasoning을 끝내려고 하면, 이를 방지하고 "Wait"이라는 추가 텍스트를 강제 입력
  • 모델이 reasoning을 계속하도록 유도하면서 자신의 답변을 검토 및 수정할 기회를 제공
  • 결과적으로, 초기 답변이 틀렸더라도 추가 reasoning을 통해 정답을 찾을 확률이 높아짐

---

Budget Forcing

테스트 시간 동안 모델의 사고 과정을 제어하기 위한 기법으로,
Maximum thinking tokens 강제 : 모델이 설정된 최대 사고 토큰 수를 초과할 경우, 모델의 사고 과정을 "Final Answer: "를 주입하여 강제로 종료하여 최종 답변을 생성
Minimum thinking tokens 강제 : 모델이 사고를 충분히 하지 않고 일찍 종료하려고 할 때, "생각 종료" 토큰의 생성을 방지하고 "기다려(Wait)"라는 단어를 추가하여 모델이 계속해서 답변을 다듬도록 유도

예제 문제: "raspberry" 단어에 포함된 'r'의 개수는 몇 개인가?
일반적인 모델 동작

• 모델이 reasoning을 수행한 후 잘못된 답(2)을 생성하고 reasoning을 멈춤

Budget Forcing 적용 후

• 모델이 reasoning을 끝내려고 할 때, "Wait" 토큰을 추가하여 reasoning을 더 수행하도록 유도함
• 결과적으로 모델이 다시 reasoning을 수행하여 정답(3)을 도출함

해석:

• "Wait"을 추가하면 모델이 자신의 reasoning을 검토할 기회를 얻음
• Budget Forcing이 적용되면, reasoning 중 실수를 스스로 수정할 확률이 높아짐

---

Baseline

저자들은 다음과 같은 여러 test-time scaling 방법으로 budget forcing을 벤치마킹하였다.

1. 조건부 길이 제어 방법: 프롬프트에서 모델에게 생성 길이를 지정하는 방식에 의존한다.
   • (a) Token-conditional control: 프롬프트에서 사고 토큰 수의 상한을 지정한다.
   • (b) Step-conditional control: 각 단계가 약 100개의 토큰으로 구성된다고 가정하고, 생각하는 단계 수의 상한을 지정한다.
   • (c) Class-conditional control: 모델에게 짧은 시간 또는 긴 시간 동안 생각하도록 지시하는 두 개의 프롬프트를 사용한다.
2. Rejection sampling: 생성된 결과가 사전 정의된 계산 예산을 충족할 때까지 샘플링을 반복한다. 이 방식은 응답의 길이에 따른 사후 확률 분포를 포착한다.

Experiments

 

다음은 s1-32B를 다른 모델들과 비교한 결과이다. (# ex.는 fine-tuning에 사용된 예제 수, BF는 budget forcing)