Skip to content

[論文閱讀] s1: Simple test-time scaling

Last Updated on 2025-03-10 by Clay

S1 論文核心內容

Test-Time Scaling 現在是拓展 LLM 性能的一個熱門方案，藉由在模型回答問題前，先經過一番『思考』整理，進而回答出更準確的答案。

s1 的論文做出了以下貢獻：

  • 整理了一個名為 s1K 的資料集並開源
  • 提出了 Budget Forcing 這一方法來控制 Test-Time Scaling 的計算量
    • 強制中止: 在思考的 Tokens 超出設定值後強制停止思考
    • 強制生成: 在模型結束思考時，插入 [Wait] Token 來強迫模型繼續思考
  • 以 Qwen-2.5-Instruct 為基底，在 16 張 H100 上進行微調（僅須 26 分鐘），訓練方式為 SFT

推理資料集建構

最早研究團隊收集了 59K 筆資料，但他們希望能夠以最少的資源、最簡單的方式達成目標（讓模型學會 Test-Time Scaling），所以他們持續清理資料集直到只剩 1K 的資料量。

他們採用以下三種指標來挑選資料：

  1. 品質（Quality）：移除了 API Error 而錯誤的訓練資料，接著移除格式錯誤的資料
  2. 難度（Difficulty）：對於 Qwen-2.5 7B 和 32B 兩個模型，並使用 Claude-3.5-Sonnet 與正確的參考答案比較，如果 7B 和 32B 都能答對代表太容易；同時也使用 Tokenizer 計算 Tokens 數量（假設越難的問題需要越多的 Think Tokens）
  3. 多樣性（Diversity）: 在分類資料多樣性的部份，使用 Claude-3.5-Sonnet 並基於 Mathematics Subject Classification (MSC) System 進行分類

Budget Forcing

論文中嘗試了兩種不同的 Test-Time Scaling 的方式，一個是序列擴展（Sequential Scaling）、一個是並行擴展（Parallel Scaling）。

  • 序列擴展：讓模型逐步生成推理過程，每一階段的推理都依賴於前一階段的結果，過程是線性發展。優點是符合人類思維結果（個人倒覺得人類並不一定完全是線性思考 XD，但至少這種形式比較適合人類理解）；缺點是可能受限於長文理解能力，誤解了先前的思考步驟。
  • 並行擴展：讓模型同時生成多個推理路徑，再進行合併與篩選。優點是速度快、並能讓模型同時思索多種可能性，但缺點是更容易產生幻覺，並對像是數學證明與邏輯推理這種需要長步驟演繹的任務不太適合。

在 S1 論文中，主要專注在 Budget Forcing，也可以從他們做的兩種不同的實驗中看出來，其在 GPQA Diamond 中，序列擴展更適合這類任務。


而下圖則是一個明顯的，透過追加 Wait Token 來強迫模型繼續思考的形式。


另外，我個人也滿好奇 Budget Forcing 到底要控制在多少 Tokens 性能才會最好呢？論文中也明確提出了 AIME24 的實驗結果。

看起來收益約在強迫思考到 4,096 ~ 8,192 這一區間達到飽和。

評估 Benchmarks

研究團隊在 3 個不同的 benchmarks 上進行測試：

  • AIME24（美國數學邀請賽）
  • MATH500（競賽數學題）
  • GPQA Diamond（博士級科學問題）

我們可以看到 S1-32B 的分數乍看之下不怎麼出彩，但實際上在 AIME24 中是贏過 GPT-o1-preview 的、並且在 MATH500 中與 GPT-o1 只差一點點的分數（當然這個 Benchmark 可能是最簡單的，畢竟大家分數都很接近），在 GPQA Diamond 中則是與 GPT-o1-mini 分數接近。

跟原始的 Qwen2.5-32B-Instruct 相比，在各個 Benchmark 的分數都成長了：

  • AIME24: 提昇 30 個百分比
  • MATH500: 提昇 9 個百分比
  • GPQA Diamond: 提昇 10.6 個百分比

確實能看出 S1K 資料集對於提昇模型在邏輯性強的任務中的幫助。

另外還有一個我很好奇的問題，就是精選過後的 1K 資料集、跟原本的 59K 資料集相比又如何呢？論文中也放出了消融實驗的結果：

可以看到 s1K 的資料集優於其他的所有變體（隨機選擇、多樣性選擇、最長選擇），而甚至在 MATH500 的任務上還以些微之差贏過了 59K-full 的版本。

可以想像成只需要 1/60 的成本，就逼進了使用全部資料集下去微調的結果。證明了 Test-time Scaling 可以在高品質的小資料集上提昇原本模型的表現。

總結

我個人認為，研究團隊是透過實驗證明了兩件事情：第一是乾淨且高品質的小資料集是能夠追上大量資料集的效果（不過資料品質跟多樣性同樣需要兼顧就是了）、第二就是 Test-time Scaling 可以透過輕量資料集讓模型掌握。

不過關於第二點，也存在著『Qwen 天生在預訓練中就見識過推理行為』，進而導致它更容易在 SFT 階段被這種輕量資料集引出推理的能力。

接下來我會繼續閱讀《Cognitive Behaviors that Enable Self-Improving Reasoners, or, Four Habits of Highly Effective STaRs》這篇論文，進一步確認是否有以上可能。

私心也想用 Gemma-2-9B 在 s1K 的資料集上測試看看，在其他的 Benchmarks 是否也會有提昇呢？還是甚至會倒退呢？

References

Read More

Tags:

Leave a Reply取消回覆

Exit mobile version