핵심 : 고품질 1000개 데이터와 reasoning budget을 조절함으로써 높은 성능에 도달 가능하다.
0. Abstract
test-time에서 추가적인 자원(GPU, 시간)을 제공해서 성능을 높이는 것이 가능하다는걸 OpenAI의 모델 o1이 보여줬지만, 그 구체적인 방법론은 공개되지 않았다. 따라서 저자들은 최대한 단순하게 test-time scaling의 성능을 달성하기 위한 방법을 제안한다. 첫 번째는 난이도, 다양성, 품질을 검증하여 선별된 1000개의 데이터셋인 s1K로 SFT한다. 두 번째는 budget forcing으로, 너무 길면 추론을 멈추고 너무 짧다면 “Wait” 토큰을 넣어서 더 길게 만든다. 이를 통해 test-time compute를 통제한다.
Qwen2.5-32B-Instruct 모델에 이 두 가지 방법을 적용한 모델인 s1-32B는 o1-preview의 성능에 도달하며, 단순히 budget forcing을 scaling함으로써 AIME24에서 성능 향상을 보인다.
1. Introduction
기존 연구들은 train-time에서 어떻게 scaling up을 하는지에 집중해왔다. 이를 바탕으로 새롭게 test-time scaling이 나왔고, OpenAI의 o1이 그 사례이다. o1을 재현하기 위해 다양한 시도들이 있었지만, Monte Carlo Tree Search, Multi-agent, DeepSeekR1의 백만개 샘플로 RL하기, 제대로 test-time scaling을 하는 시도는 없었다.
따라서 해당 연구에서 1000개의 샘플로만 SFT하고 thinking duration을 통제하는 budget forcing만으로 test-time copmute를 구현한다. Gemini Thinking Experimental로 distill 했으며, H100 GPU 16개로 26분만 학습한다. budget forcing은 EoT 토큰을 지우고 “Wait”을 넣거나, 강제로 EoT 토큰을 넣는 것으로 구현한다.
데이터셋을 랜덤하게 만들거나 가장 긴 trace만 고르는건 오히려 성능이 안좋아져서 s1K의 중요성을 보여주고, 다른 test time compute와 비교해서 budget-forcing이 가장 control이 좋고 성능 향상이 있음을 보인다.
2. Reasoning data curation to create s1K
2.1 Initial collection of 59K samples
먼저 큰 데이터셋을 만드는데, quality, difficulty, diversity를 고려해서 수집한다. 기존에 존재하는 데이터셋들에서 수학 데이터셋을 많이 사용하였고, 다양한 도메인을 다루는 OlympicArena 등을 사용하였다. 보완하기 위해서 어려운 수학문제들로 구성된 2개의 데이터셋을 새롭게 만든다.
답변을 생성할 때는 Google Gemini Flash Thinking API를 사용하여 reasoning trace와 response를 만든다. 59K개의 question-reasoning-trace-solution triplet이 생기고, 평가용 데이터셋과 비슷한 데이터는 8-gram을 사용하여 제거한다.
2.2 Final Selection of 1K samples
59K를 전체 사용할 수 있으나, 가장 단순한 접근을 찾기 위해 1000개의 샘플로만 데이터셋을 만들기 위한 3가지 필터링 단계를 사용한다.
Quality
API 에러가 발생한 샘플들을 먼저 지우고, 포맷팅과 관련된 오류(아스키 아트, 잘못된 이미지 참조, 넘버링 오류)가 있는 샘플들을 제거한다. 해당 단계에서 384개의 고품질이자 필터링이 더 필요없는 샘플들을 우선적으로 특정한다.
Difficulty
난이도는 모델 성능과 답변 길이를 바탕으로 측정한다. 답변 길이를 바탕으로 질문들의 난이도를 예측하고, 2개의 다른 크기를 가진 모델 중 하나라도 맞는 샘플들은 제거한다. 2개의 모델을 사용함으로써 하나의 모델만 사용했을 때 실수로 너무 쉬운 샘플도 통과되는 것을 방지한다.
Diversity
Claude 3.5 Sonnet으로 질문들을 MSC (Mathematics Subject Classification) 체계로 분류한다. 하나의 도메인을 랜덤하게 고르고 가장 긴 reasoning trace를 가지고 있는 샘플을 선택한다. 1000개가 만들어질 때까지 각 도메인에서 이 과정을 반복한다.
3. Test-time scaling
3.1 Method
test-time scaling을 2가지 유형으로 분류하는데, 첫 번째는 Sequential하게 초기 연산에 후기 연산이 의존하는 것이고, 두 번째는 연산이 병렬적으로 이루어지는 (i.e. majority voting) 경우이다. 전자가 연산을 쌓아올리기 때문에 scaling하기에 더 용이하다고 이야기하며, 전자에 초점을 맞춘다.
Budget forcing
decoding-time에서 간단한 intervention을 통해 thinking token의 최대/최소를 강제한다. 최대를 맞추기 위해서는 너무 길어질 때는 EoT 토큰 delimiter와 “Final Answer”를 추가해서 추론 단계를 멈추고 현재 상태에서 답변을 하도록 한다. 최소를 맞추기 위해서는 EoT를 지우고 “Wait” 토큰ㅇ르 생성하게 해서 더 길게 reasoning trace를 만들도록 한다.
Baseline
비교 대상은 (I) Conditional length-control 방법들로 전체 프롬프트에서 토큰을 통제하거나, 스텝별로 통제하거나, 두 개의 프롬프트로 하는 것과 (II) Rejection sampling으로 원하는 budget에 맞는 샘플이 나올 때까지 뽑을 수 있다.
3.2 Metrics
정확도 뿐만이 아니라, 통제가 얼마나 가능한 지와 test-time scaling의 정도를 확인한다.
- Control: 원하는 토큰의 minimum과 maximum 사이에 들어올 수 있게 하는 지를 확인
- Scaling: 사용하는 토큰 수를 늘릴수록 성능이 올라가는 지를 확인
- Performance: 토큰 수를 최대로 늘렸을 때 달성하는 최고 성능 (현실적으로 control이나 context window 제한으로 인해 flatten된다)
4. Results
4.1 Setup
Qwen2.5-32B-Instruct 모델을 파인튜닝했으며, H100 GPU 16개로 학습했다. 평가는 AIME24, MATH500, GPQA Diamond 등의 데이터셋으로 진행했다. 비교 모델로는 OpenAI o1, DeepSeek r1, QwQ-32B-preview, Gemini2.0 Flash Thinking Experimental 등을 사용했다.
4.2 Performance
budget forcing을 사용하면 어느 정도 향상이 있지만, reasoning을 너무 늘리면 어느정도 flattern되기도 하고 더 나은 추론을 하기보다 반복적인 loop에 들어가는 현상도 발생한다. 효율성에 비교하면 s1-32B 모델이 가장 sample-efficient함을 알 수 있다. AIME2024에서는 Gemini의 성능을 거의 따라잡으면서 distillation 과정이 효과적임을 알 수 있다.
5. Ablations
5.1 Data Quantity, Diversity, and Difficulty
- Only Diversity: 난이도와 상관없이 도메인만 다양하게 되도록 1K를 뽑으면 성능이 낮아졌다. (랜덤하게 뽑는 것과 유사한 성능)
- Only Difficulty: 가장 긴 reasoning trace를 가진 샘플들로만 구성했을 때, GPQA 성능은 좋아졌지만 s1K에 비하면 낮았다.
- Maximize Quantity: 59K 전체를 사용했을 때는 성능이 유사하면서 조금 높았지만, s1K를 사용한 것의 효율성이 더 좋다는 것을 이야기한다.
5.2 Test-time scaling methods
Budget forcing을 token/step/class-conditional한 control과 Rejection sampling과 비교한다. 가장 성능이 좋은 방법은 Budget forcing이었으며, token-conditional control은 토큰 수를 제대로 세지 못하는 오류를 계속 범했다. step-conditional control은 step 수와는 무관하게 전체 토큰 수가 일정하게 유지되어, 모델이 goal에 대해 hacking함이 드러난다. Class-conditional control은 test-time compute와 성능이 잘 scaling됨을 확인한다. Rejection sampling의 경우, 오히려 토큰이 늘어날 수록 성능이 줄어드는데, 긴 response는 이미 실수를 해서 backtracking을 시도하는 경우가 많기 때문이라고 이야기한다.
6. Discussions and related work
6.1 Sample-efficient reasoning
여러 모델들이 RL이나 많은 데이터로 SFT하면서 o1의 성능을 따라가려는 시도를 했다. 이 연구에서는 1000개의 샘플들로만 SFT를 하는게 유사한 성능을 가질 수 있음을 보인다. 해당 저자들은 이미 많은 양의 추론과 관련된 토큰들로 학습되어서 이미 추론을 할 능력 자체는 내재되어 있다고 추측한다. sample-efficient한 training을 통해 이를 활성화시키고, budget forcing으로 scaling한다고 주장한다. 이는 LIMA 논문에서 1000개 예시만으로 사용자 선호도로 align시킬 수 있다는 주장과 유사하다.
모델의 성능을 끝까지 밀어붙여 평가하기 위한 벤치마크들도 많아졌다. 기존 연구들은 특별한 도메인들에 대해 corpora를 구성하거나, 합성 데이터를 사용한다. 추론 능력을 높이기 위해서는 추가적인 학습이나, 프롬프팅과 같은 방법론들도 시도되고 있다.
6.2 Test-time scaling
앞서 말한 것처럼 1) parallel과 2) sequential 유형으로 구분된다. 전자는 majority voting이나 Best-of-N 등이 포함된다. 후자는 앞선 시도들을 바탕으로 refine하는 방법이 된다. Tree-based search 방법은 하이브리드 접근이며, REBASE와 같은 방법은 PRM으로 exploitation과 pruning의 균형을 맞춘다. Reward 모델에서도 outcome based는 parallel에 도움이 되고, process based는 sequential에 도움이 된다.
test-time scaling의 한계로는 결국에는 flatten out된다는 점과, context window의 제한이 있다는 것이다. 추가적인 연구 방향으로는 reasoning을 연장할 때 더 다양한 토큰을 사용하거나, RL 모델들의 budget forcing 효과 등이 있다., sequential scaling의 문제점을 parallel scaling으로 일부 보완할 수 있다는 것도 발견한다.
EMNLP 2025; ICLR 2025 Reasoning and Planning Workshop Best Paper Award
논문 링크:
https://arxiv.org/abs/2501.19393v2