1KのデータセットでSFTするだけでもo1-previewを超えられる

paper: s1: Simple test-time scaling, Feb 2025

https://arxiv.org/abs/2501.19393

Summary

• 慎重にキュレートした1KのデータセットでQwen2.5-32B-InstructをSFTするとAIMEなど数学のベンチマークでo1-previewのスコアを上回った。
• 良いデータセットを選ぶ指針: 1) Quality, 2) Difficulty, 3) Diversityがどれも重要。どれかが欠けると劣った性能を示した。
• 推論時スケーリングをコントロールする手法を提案。Budget Forcing (BF)

Method

キュレート方法

最初にAIME, NuminaMath, OlympicArenaなど複数のデータセットから59Kのサンプルプールを用意し、Google Gemini Flash Thinking APIを使って思考過程を生成する。

1. Quality: APIエラーが出たものをのぞいて54,116サンプルを得る。そこからアスキーアートや存在しない画像のURLなど余計なノイズを含むものを除いて51,581のサンプルを得る。
2. Difficulty: 1) model performance, 2) reasoning path lengthの2つの指標でフィルタリングする。Qwen2.5-7B-Instruct, Qwen2.5-32B-Instructを使って正解可能なeasyサンプルを除き、最終的に24,496サンプルを得る。
3. Diversity: American Mathematical Societyが規定する50の問題のカテゴリ(Geometry, Number theoryなど)をClaude 3.5 Sonnetを使って分類させ、そのうちからreasoning path lengthが最も長いものを1問ずつサンプリングする。この処理を繰り返し、最終的に1,000サンプルをえる。これをs1Kと呼称する。

Budget Forcing(BF)とは？

• think tokenの長さを特定の範囲に収まるように推論時介入を行う。
  • 短すぎる場合: </think>の確率を下げて"Wait"というトークンを追加。
  • 長すぎる場合: 強制的に思考を打ち切り、"Final Answer:"というトークンを追加。

実験結果

• s1Kを使ってQwen2.5B-32B-InstructをSFT。この時点でAIME, MATHでo1-previewのスコアを超える。
• さらに、Budget Forcingを適用して推論時スケーリングを行うことで追加の改善を確認。
• ablationにて59Kのfull datasetは59倍もサンプル数が多いにも関わらず、それを使ってSFTしたモデルはs1Kとわずかなゲインしか得られなかった(AIME: 50.0->53.3, MATH: 93.0->92.8)。→s1Kは極めてサンプル効率が高い
• Budget Forcingの代わりにrejection sampling(目的の長さの思考過程が得られるまでサンプリングを繰り返す)を行った場合、推論時スケーリングは逆の傾向を示した（トークン長が長いほど性能が悪化）。→問題の難しさに応じて最適な思考過程の長さが決まっている？

考察

• なぜ1Kデータセットでo1-previewを越えられたのか？ →著者らの仮説: 「pretraining時にすでに推論方法は学習しており、SFTでそれが活性化され、Budget Forcingでさらにスケールされた」

所感

• この論文は、推論時スケールの効果をbudgetと引き換えにコントロールする手段を提供したという意味で意義がある。
• また、DeepSeek-R1は具体的な学習データセットを公開していないが、本論文により、o1-previewと同等以上の性能のモデルを学習するデータのレシピが公開されたということに意義がある。
• キュレート方法に関する知見は、今後の研究や応用において助けになるだろう。
• わずか1KのSFTでもo1-previewを超えるスコアを出すのに十分、という事実は、「推論能力の獲得に事後学習がどの程度寄与しているか？」という疑問を提起する。