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

안녕하세요. 앨런입니다.

오늘은 s1: Simple test-time scaling 논문을 리뷰하려고합니다.

 

0. 요약

• OpenAI의 o1 모델이 테스트 시간 확대 기법을 성공적으로 구현했으나 그 방법론이 공개되지 않았음
• 이를 구현하기 위해 s1K 데이터셋 을 구축하였고 Test time scaling기법으로 추론 모델의 성능을 높였음
• 32B 급 모델로 노트북에서도 구동 가능한 학습 방법론을 제시 했다는 점에서 의의가 있음
• 16개의 NVIDIA H100 GPU에서 26분을 소요하여 적은 비용으로 가볍고 우수한 성능의 모델을 만들었다는 점에서 논문의 의의가 있음

1. 본 논문의 목적

• OpenAI의 o1 모델이 테스트 시간 확대 기법을 성공적으로 구현했으나 그 방법론이 공개되지 않았음
• 본 연구는 가장 간단한 방법으로 Test time scaling (테스트 시간 확대) 로 강력한 추론 성능을 달성하고자 함

• 그림에서 X축 (Average thinking time (tokens)): 모델이 문제를 해결하는 데 평균적으로 소비하는 "생각의 시간", 즉 사용하는 토큰의 수를 나타냄
• Y축 (Accuracy (%)): 모델이 문제를 올바르게 해결한 비율인 정확도를 백분율로 나타냄
• MATH500, AIME24, GPQA Diamond 데이터셋에서 생각이 길어질수록 정확도가 향상되는 것을 볼 수 있음

2. 제안 방법

2.1 s1K 데이터셋 구축

• 1,000개의 질문과 추론 과정을 포함하는 s1K라는 소규모 데이터셋을 구축함
• NuminaMATH, OmniMATH 등 16개 데이터셋의 59K 샘플로부터 최소한의 자원으로 1,000개의 샘플을 선택
• 데이터 선정은 품질(quality), 난이도(difficulty), 다양성(diversity)의 세 가지 기준을 기반으로 샘플을 선정
• 1단계: 질 (Quality)
  • API 오류로 인해 불량 샘플 제거: 54,116 샘플로 축소
  • 형식 문제 (ASCII 아트, 이미지 참조 오류 등)로 저품질 샘플 식별: 51,581 샘플로 감소
• 2단계: 난이도 (Difficulty)
  • 두 모델(Qwen2.5-7B-Instruct 및 Qwen2.5-32B-Instruct)을 사용하여 각 질문 평가
  • 모델이 올바르게 해결한 문제는 너무 쉬운 것으로 간주하며 제거
  • 이로 인해 24,496 샘플로 축소
• 3단계: 다양성 (Diversity)
  • 질문을 특정 도메인으로 분류 (예: 기하학, 생물학 등).
  • 각 도메인에서 고른 분포에 따라 문제를 무작위로 샘플링하여 최종 1,000개 샘플을 선택.

 

 

• 위 그림에서 (왼쪽) s1K 데이터셋의 분포를 보여주며 (오른쪽) 은 학습 데이터의 수에 따른 MATH500 의 정확도를 보여줌
• s1 모델은 적은 수의 예제 (학습데이터) 로 높은 정확도를 보여주고 있음

2.2 Test time Scaling 기법

• Test time scaling 기법의 종류
  • 순차적 방법 (Sequential): 후속 계산이 이전 계산의 결과에 의존함. 즉, 각 단계의 결과가 다음 단계의 입력으로 사용되어 깊은 추론과 iterative refinement(반복적인 개선)를 가능하게 한다. 예를 들어, 모델이 문제를 해결하기 위해 여러 단계의 사고를 거치는 경우가 해당된다.
  • 병렬적 방법 (Parallel): 계산이 독립적으로 수행됨. 각 계산 단계가 서로의 결과에 영향을 미치지 않으며 동시에 진행될 수 있다. 이는 주로 다수의 결과를 동시에 생성하거나 비교할 때 사용된다. 예를 들어, 다수의 샘플을 동시에 처리하여 결과를 집계하는 경우가 있다.
• 예산 강제 (Budget Forcing) :
  • 테스트 시간 동안 모델의 사고 과정을 제어하기 위한 기법으로,
    Maximum thinking tokens 강제 : 모델이 설정된 최대 사고 토큰 수를 초과할 경우, 모델의 사고 과정을 "Final Answer: "를 주입하여 강제로 종료하여 최종 답변을 생성
    Minimum thinking tokens 강제 : 모델이 사고를 충분히 하지 않고 일찍 종료하려고 할 때, "생각 종료" 토큰의 생성을 방지하고 "기다려(Wait)"라는 단어를 추가하여 모델이 계속해서 답변을 다듬도록 유도

3. 실험 결과

• Qwen 2.5 32B-instruct 모델을 위 데이터셋으로 supervised finetuing을 통해 s1-32B 모델을 구축
• 훈련 후, s1-32B 모델은 경쟁 수학 문제 ((MATH and AIME24) 에서 o1-preview보다 최대 27% 더 높은 성능을 보였음
• 모델 학습 및 평가 설정
  • 모델 학습: Qwen2.5-32B-Instruct 모델을 사용하여 s1K 데이터셋으로 학습함
  • 평가 벤치마크: AIME24, MATH500, GPQA Diamond의 세 가지 대표적인 추론 기준으로 평가함

 

• 성능
  • AIME24 성능: s1-32B 모델은 AIME24에서 평균 56.7%의 정확도
  • MATH500 성능: MATH500에서는 93.0%의 정확도
  • GPQA Diamond 성능: GPQA Diamond에서는 59.6%의 정확도
• 테스트 시간 스케일링
  • 예산 강제(Budget Forcing): 모델의 추론 중에 더 많은 계산을 하도록 설정하여 성능을 개선
  • 성능 향상: Budget Forcing을 통해 AIME24 성능을 개선할 수 있었으나, 특정 시점 이후에는 성능이 정체되었음

4. 한계점

• 예측 성능 한계: 예측 성능을 높이기 위해 '예산 강제(budget forcing)' 기술을 사용하였지만, 이 방법에는 한계가 있음.
• 평탄화 현상: 성능 개선이 한계에 도달하면 더 이상 향상되지 않는 평탄화 현상이 있음.
• 컨텍스트 윈도우의 제한: 기반이 되는 언어 모델의 컨텍스트 윈도우가 성능 향상을 제약함. 이는 모델이 한 번에 처리할 수 있는 빈 공간을 말함.
• 추가적인 연구 필요: 계속해서 테스트 시간에서 성능을 개선하려면 새로운 방법이 필요함. 예를 들어, 다양한 문자열을 통해 예산 강제를 개선하거나, 강화학습을 통해 새로운 테스트 시간 확장 방법을 탐구할 수 있음.

참고문헌

• Paper : https://arxiv.org/abs/2501.19393

s1: Simple test-time scaling

 

• Github : https://github.com/simplescaling/s1

GitHub - simplescaling/s1: s1: Simple test-time scaling

 

• huggingface : https://huggingface.co/simplescaling/s1-32B

simplescaling/s1-32B · Hugging Face