文｜季逸超 Peak，Peak Labs 創(chuàng)始人，真格基金 EIR（Entrepreneur In Residence，入駐創(chuàng)業(yè)者）

從 OpenAI o1 發(fā)布以來(lái)，我就把復(fù)現(xiàn) o1 作為業(yè)余時(shí)間的興趣項(xiàng)目，其過(guò)程中獲得了很多有趣的知識(shí)與洞察 (與失敗)。考慮到個(gè)人能力、精力、財(cái)力有限，我很可能不是走到最后的選手，因此我覺(jué)得很有必要將這些價(jià)值數(shù)十張 H100 的經(jīng)驗(yàn)及時(shí)分享出來(lái)。

本文的英文版同步發(fā)表于：

https://medium.com/@peakji/a-small-step-towards-reproducing-openai-o1-b9a756a00855

01

TL;DR

Steiner 是一個(gè) reasoning 模型，能在推理時(shí)以自回歸的形式探索多種路徑，并在必要時(shí)自主進(jìn)行驗(yàn)證或回溯。訓(xùn)練過(guò)程分為三個(gè)步驟：首先，通過(guò)隨機(jī)截?cái)嗪铣傻?reasoning path 并進(jìn)行引導(dǎo)式再補(bǔ)全，將每個(gè)樣本轉(zhuǎn)化為從問(wèn)題到答案的有向無(wú)環(huán)圖（DAG）。接著，從 10K 個(gè) DAG 中采樣出 50K 個(gè)具有回溯節(jié)點(diǎn)的 reasoning path 樣本。最后，使用 DAG 中節(jié)點(diǎn)的入邊和出邊數(shù)量及其距離作為 heuristic rewards，通過(guò)強(qiáng)化學(xué)習(xí)讓模型學(xué)會(huì)平衡探索的廣度和深度。該模型在 GPQA-Diamond 數(shù)據(jù)集上實(shí)現(xiàn)了 +5.56 的提升，但未能復(fù)現(xiàn) inference-time scaling。

模型下載地址：

https://huggingface.co/collections/peakji/steiner-preview-6712c6987110ce932a44e9a6

02

簡(jiǎn)介

Steiner 是一系列使用強(qiáng)化學(xué)習(xí)在合成數(shù)據(jù)上訓(xùn)練得到的 reasoning 模型，能夠在推理時(shí)以自回歸的形式嘗試多種推理路徑，并在必要時(shí)自主進(jìn)行驗(yàn)證和回溯，從而實(shí)現(xiàn)在同一上下文中對(duì)隱含的搜索樹(shù)的線性遍歷。

Steiner 能夠在推理時(shí)以自回歸的形式嘗試多種推理路徑，并在必要時(shí)自主進(jìn)行驗(yàn)證和回溯，從而實(shí)現(xiàn)在同一上下文中對(duì)隱含的搜索樹(shù)的線性遍歷。

Steiner 是 季逸超 (Yichao 'Peak' Ji) 的個(gè)人興趣項(xiàng)目，受 OpenAI o1 啟發(fā)，最終目標(biāo)是復(fù)現(xiàn) o1 并驗(yàn)證 inference-time scaling 曲線。Steiner-preview 模型是一個(gè) work-in-progress。開(kāi)源的目的是因?yàn)槲矣^察到以選擇題為主要構(gòu)成的自動(dòng)化評(píng)估手段較難充分體現(xiàn) reasoning 模型的進(jìn)展：事實(shí)上，「正確答案一定在選項(xiàng)中」這一前提就與真實(shí)的 reasoning 場(chǎng)景并不對(duì)齊，它會(huì)鼓勵(lì)模型進(jìn)行代入驗(yàn)證而非開(kāi)放式探索。因此，我選擇將階段性成果開(kāi)源出來(lái)，并在精力允許的情況下 build in public，從而在分享知識(shí)的同時(shí)獲得更多來(lái)自真實(shí)人類的評(píng)估和反饋。

??Disclaimer: 到目前為止，Steiner 可以在無(wú)需 Chain of Thought (CoT) prompting 和 agent framework 的情況下獲得較高質(zhì)量的 zero-shot 結(jié)果，但是仍未能復(fù)現(xiàn) o1 展示出的 inference-time scaling 能力：在使用特殊的 logits processor（https://gist.github.com/peakji/f81c032b6c24b358054ed763c426a46f）對(duì) reasoning tokens 進(jìn)行干預(yù)的實(shí)驗(yàn)中，增加額外的 reasoning steps 沒(méi)能提升性能，反而在 MMLU-Pro 和 GPQA 等 benchmark 上有所下降。因此 Steiner 暫時(shí)還不能被視為 OpenAI o1 的成功復(fù)現(xiàn)，可能在訓(xùn)練方法和數(shù)據(jù)質(zhì)量上都存在不足，請(qǐng)謹(jǐn)慎參考。

03

背景

與傳統(tǒng) LLM 相比，OpenAI o1 最顯著的變化是在推理時(shí)引入了 reasoning tokens，進(jìn)而實(shí)現(xiàn)了 inference-time scaling：即通過(guò)增加推理時(shí)的 compute budget 來(lái)提升模型的表現(xiàn)。當(dāng)我們討論 inference-time scaling 時(shí)，最直覺(jué)的方法可能是引入 tree search 或 agentic framework。但是，我在閱讀 o1 相關(guān)的 (有限的) 官方信息時(shí)，注意到其 report 的 benchmark 多是 pass@1 和 majority voting，而且 OpenAI 團(tuán)隊(duì)也提到 o1 是 single model 而非 system。這讓我對(duì)其實(shí)現(xiàn)的方法非常好奇。

OpenAI API 中至今仍未將 reasoning tokens 的具體內(nèi)容開(kāi)放給開(kāi)發(fā)者，但好在 tokens 用量的統(tǒng)計(jì)數(shù)據(jù)中會(huì)包含 reasoning tokens 的數(shù)量 (畢竟要根據(jù)這個(gè)向開(kāi)發(fā)者收錢)。借此我設(shè)計(jì)了一個(gè)簡(jiǎn)單的實(shí)驗(yàn)，嘗試通過(guò) o1 的 API 來(lái)統(tǒng)計(jì) completion (包含 reasoning) tokens 的數(shù)量與總請(qǐng)求時(shí)長(zhǎng)的關(guān)系。我們知道，若使用 tree search，為了復(fù)用 cache 和提升 GPU 利用率會(huì)盡可能的并行推理，那么得到曲線應(yīng)是 sub-linear 的。然而，實(shí)驗(yàn)的結(jié)果卻是幾條漂亮的直線，o1-mini 的波動(dòng)甚至比 GPT-4o-mini 還小：

OpenAI o1 可能仍是一個(gè)進(jìn)行線性自回歸解碼的模型，但這不代表它在推理階段沒(méi)有進(jìn)行 「搜索」。

上述實(shí)驗(yàn)讓我產(chǎn)生了一種猜想：OpenAI o1 可能仍是一個(gè)進(jìn)行線性自回歸解碼的模型，但這不代表它在推理階段沒(méi)有進(jìn)行「搜索」。想象有一顆搜索樹(shù)，我們?cè)趯?duì)它進(jìn)行遍歷的時(shí)候，產(chǎn)生的軌跡其實(shí)是線性的。如果我們能夠訓(xùn)練一個(gè)自回歸語(yǔ)言模型，讓它不僅能生成推理路徑，還能在必要時(shí)進(jìn)行驗(yàn)證、回溯、或切換思路，那么在最理想的情況下它實(shí)際上完成了在同一上下文中對(duì)隱含的搜索樹(shù)的線性遍歷。這種線性遍歷聽(tīng)起來(lái)很浪費(fèi)，但與并行搜索相比，它具有三個(gè)極其優(yōu)秀的性質(zhì)：

1. 前置的嘗試無(wú)論對(duì)錯(cuò)都在 context memory 中，每一步?jīng)Q策均基于過(guò)去的完整信息；

2. 隱含的 backtracking 無(wú)需保證目標(biāo)節(jié)點(diǎn)已經(jīng)在搜索樹(shù)內(nèi)，探索更加自由；

3. 工程層面可以復(fù)用一切現(xiàn)有的經(jīng)過(guò)高度優(yōu)化的推理 infrastructure。

想象有一顆搜索樹(shù)，我們?cè)趯?duì)它進(jìn)行遍歷的時(shí)候，產(chǎn)生的軌跡其實(shí)是線性的。

04

方法

要獲得具有線性搜索能力的模型并不容易，無(wú)論是數(shù)據(jù)合成還是模型訓(xùn)練都面臨許多挑戰(zhàn)。

首先，目前可以取得的 reasoning 數(shù)據(jù)集大多僅有合成的 CoT 或 reasoning steps，它們一般是通過(guò)將「問(wèn)題-答案」tuple 輸入給強(qiáng)大的 LLM，然后要求模型拆解思考過(guò)程獲得的。這種方法決定了這些數(shù)據(jù)集不會(huì)包含合理的 backtracking 節(jié)點(diǎn)，故而利用這種數(shù)據(jù)訓(xùn)練得到模型其實(shí)只學(xué)到了 shortcut，或者說(shuō)是將 CoT internalized 了。

針對(duì)這個(gè)問(wèn)題，我設(shè)計(jì)了兩種數(shù)據(jù)合成與擴(kuò)增方法：

1. 對(duì)上述 shortcut 數(shù)據(jù)集進(jìn)行隨機(jī)截?cái)啵㈦[藏正確答案，讓強(qiáng)大的 LLM 基于被截?cái)嗪蟮那熬Y去嘗試正向推理一定的 steps，然后再提供正確答案，以獲得 backtracking 樣本；

2. 對(duì)上一步產(chǎn)生的 steps 進(jìn)行聚類后賦予唯一 ID，將同一個(gè)問(wèn)題下的所有 steps 構(gòu)建為有向無(wú)環(huán)圖 (DAG)，并對(duì) DAG 進(jìn)行隨機(jī)采樣以獲得多項(xiàng)式量級(jí)的 reasoning path 樣本。

3. 通過(guò)上述方法 (以及大量人工和奇技淫巧)，我最終獲得了 10K 張有向無(wú)環(huán)圖，并在此基礎(chǔ)上采樣了 50K 條帶有 backtracking 的 reasoning path 樣本。其中每個(gè)樣本的 reasoning tokens 的平均數(shù)量約 1600 個(gè)，剛好與先前對(duì) o1/o1-mini 進(jìn)行測(cè)試時(shí)的統(tǒng)計(jì)數(shù)據(jù)非常接近！考慮到訓(xùn)練成本，我僅保留了 reasoning tokens 數(shù)量在 4096 以下，且 prompt + reasoning + completion 總 tokens 數(shù)量在 8192 以下的樣本。

接著，我將 Steiner 系列模型的訓(xùn)練分為三個(gè)階段：

1. Continual Pre-Training (CPT)：混合普通文本語(yǔ)料和 reasoning path 進(jìn)行訓(xùn)練，目的是讓模型習(xí)慣 long reasoning output，并且初步訓(xùn)練新增的 14 種特殊 tokens（https://huggingface.co/peakji/steiner-32b-preview/blob/278989ea17f74b14e2b32d9544eb53a17b4ad087/special_tokens_map.json#L16-L29）的 embeddings。需要指出的是，在一些小參數(shù)量模型的測(cè)試中表明，這一步可能是冗余的，直接在 SFT 階段用大量 reasoning 數(shù)據(jù)訓(xùn)練似乎也能獲得不錯(cuò)的表征，但 32B 的 CPT 很早就做完了，我就沿用到現(xiàn)在了；

2. Supervised Fine-Tuning (SFT)：使用 chat template 進(jìn)行訓(xùn)練，目的是讓模型學(xué)會(huì)模仿 reasoning 的格式：先給每個(gè) step 想一個(gè) name，再輸出完整的 thought，然后將 thought 總結(jié)為 summary，接著對(duì)到此為止的推理進(jìn)行 reflection，最后決定下一步是 proceed、backtrack 或是結(jié)束 reasoning 并開(kāi)始正式回答問(wèn)題。你可能會(huì)好奇，開(kāi)源模型既然無(wú)需隱藏 thoughts 為什么也要像 o1 那樣生成 summary 呢？這是因?yàn)槲以跒楹罄m(xù)的具有多輪對(duì)話能力的 Steiner 模型做準(zhǔn)備。理論上來(lái)說(shuō)，經(jīng)過(guò)訓(xùn)練后可以選擇將先前對(duì)話中的完整 thought 用 summary 取代，來(lái)減少無(wú)法命中 prefix cache 時(shí)的 pre-fill 開(kāi)銷。目前 Steiner 還未針對(duì)多輪對(duì)話做優(yōu)化，僅保留 summary 可能會(huì)產(chǎn)生負(fù)向的 few-shot 效果；

3. Reinforcement Learning with Step-Level Reward (RL)：經(jīng)過(guò)前兩階段，模型已經(jīng)學(xué)會(huì)生成和補(bǔ)完 reasoning path，但它還不知道哪些是正確且高效的選擇。如果我們盲目 reward 更短的 reasoning path，則會(huì)降級(jí)為 shortcut 學(xué)習(xí)來(lái) internalize CoT。這里我設(shè)計(jì)了一種啟發(fā)式的 reward 機(jī)制：基于 DAG 中每個(gè)節(jié)點(diǎn)的入邊數(shù)量 e_i、出邊數(shù)量 e_o、距離原始問(wèn)題的距離 d_s、距離正確答案的距離 d_e 來(lái)加權(quán)，為每個(gè) step 以及整個(gè) reasoning path 賦予 reward，以引導(dǎo)模型學(xué)會(huì)平衡探索的廣度和深度。

以上思路看起來(lái)很簡(jiǎn)單，但在過(guò)去的一個(gè)月多里，每個(gè)周末（和沒(méi)加班的夜晚）我都在與 OOM 和 reward hacking 做斗爭(zhēng)。終于，在 OpenAI o1 發(fā)布后的第 38 天，訓(xùn)練出了不算太丟人的階段性成果。

05

評(píng)估

不同訓(xùn)練階段的 Steiner 模型在 GPQA-Diamond 數(shù)據(jù)集上的表現(xiàn)。Steiner 系列模型均不使用 CoT prompting。

圖中展示的是不同訓(xùn)練階段的 Steiner 模型在 GPQA-Diamond 數(shù)據(jù)集上的表現(xiàn)。可以看到引入強(qiáng)化學(xué)習(xí)階段讓模型有了 +3.53 的提升。若配合用于約束推理步數(shù)的 logits processor（https://gist.github.com/peakji/f81c032b6c24b358054ed763c426a46f），在最優(yōu)配置下可以獲得 +5.56 的提升。

使用特殊 logits processor 的 Steiner-32B 模型在 GPQA-Diamond 的各 subdomain 上的表現(xiàn)。

選擇展示該 benchmark 首先是因?yàn)?o1/o1-mini 在該數(shù)據(jù)集上有較大提升，同時(shí)該數(shù)據(jù)集的 contamination 情況較為理想。其次，也是因?yàn)槲野l(fā)現(xiàn) Steiner 在 MMLU 等數(shù)據(jù)集上與 baseline 相比沒(méi)有明顯的差異，這與 OpenAI 在關(guān)于 o1-mini 的博客（https://openai.com/index/openai-o1-mini-advancing-cost-efficient-reasoning/）中的觀察類似，可能是受限于 32B 的模型在 pre-training 階段所獲得的 world knowledge。

無(wú)論從 benchmark 還是實(shí)際使用來(lái)看，都必須承認(rèn)當(dāng)前模型的能力較 o1-mini 和 o1 還有很大差距。但另一方面，主要由選擇題構(gòu)成的自動(dòng)化評(píng)估 benchmark 可能無(wú)法充分體現(xiàn)出 reasoning 模型的能力：在訓(xùn)練階段，reasoning 模型被鼓勵(lì)對(duì)問(wèn)題進(jìn)行開(kāi)放性的探索，而選擇題存在「正確答案一定在選項(xiàng)中」這一前提，很明顯將選項(xiàng)逐一代入驗(yàn)證才是更高效的方法。事實(shí)上，現(xiàn)有的大語(yǔ)言模型均有意或無(wú)意地掌握了這種技巧，無(wú)論是否使用特殊的 prompt。當(dāng)然，正是因?yàn)樽詣?dòng)化評(píng)估與真實(shí) reasoning 需求的不對(duì)齊，讓我認(rèn)為更有必要將模型開(kāi)源出來(lái)供真實(shí)的人類進(jìn)行評(píng)估和反饋。

更大的問(wèn)題在于 inference-time scaling 實(shí)驗(yàn)：很不幸，目前 Steiner 在使用 logits processor（https://gist.github.com/peakji/f81c032b6c24b358054ed763c426a46f）增加額外的推理步數(shù)后，在幾乎所有測(cè)試的 benchmark 上都沒(méi)有提升，甚至出現(xiàn)了下降。目前猜測(cè)的原因非常多，包括但不限于：

• CPT/SFT 并未讓模型真正掌握超長(zhǎng)的 output tokens；

• RL 階段啟發(fā)式的 reward 設(shè)計(jì)未能讓模型在推理的廣度和深度之間取得平衡；

• backtrack 時(shí)受 context 影響容易確認(rèn)錯(cuò)誤答案或否認(rèn)原本正確的答案；

• 強(qiáng)行修改 logits 來(lái)增加推理步數(shù)的方法過(guò)于粗暴；

總之，這將是我接下來(lái)工作的重點(diǎn)。

06

限制

目前 Steiner 的 post-training 數(shù)據(jù)中不包含多輪對(duì)話的樣本。當(dāng)前性能最好的 Steiner 模型版本 (基于 Qwen2.5-32B) 完全不具備多輪對(duì)話能力。開(kāi)源的 Steiner-preivew 模型 (基于 Qwen2.5-32B-Instruct) 可以兼容 chat 格式，但依然不建議用于多輪對(duì)話。

與 OpenAI o1-2024-09-12 類似，Steiner 也不建議使用自定義的 system prompt 或?qū)?temperature 等 sampling 參數(shù)進(jìn)行修改。目前 Steiner 還未在多樣化的 system prompt 上進(jìn)行訓(xùn)練，而修改其他參數(shù)可能導(dǎo)致 reasoning tokens 的格式出現(xiàn)錯(cuò)誤。

Steiner 的 post-training 數(shù)據(jù)中語(yǔ)言的占比約為 90% 英文與 10% 的中文，但在 reasoning path 數(shù)據(jù)擴(kuò)增過(guò)程中幾乎完全使用英文。因此雖然模型的最終回答具備一定程度的語(yǔ)言跟隨能力，但 reasoning tokens 部分可能均為英文。

07

彩蛋

關(guān)于 Steiner 這個(gè)名字的來(lái)源，不妨問(wèn)問(wèn)模型吧：

你的名字的來(lái)源是什么？

我的名字叫做 Steiner，來(lái)源于《命運(yùn)石之門》中主角的能力「Reading Steiner」。這種能力使得主角能夠保留多個(gè)世界線的記憶，這非常契合我作為一款具有推理能力模型的特點(diǎn)：在進(jìn)行推理時(shí)，我可以嘗試多種路徑，并在必要時(shí)進(jìn)行跳轉(zhuǎn)或回溯，同時(shí)保留所有探索過(guò)的路徑的記憶。

08

引用

如果你覺(jué)得我的工作有幫助，請(qǐng)考慮在你的研究或項(xiàng)目中引用它。你的認(rèn)可將不勝感激！

@misc{ji2024steiner,
   title = {A Small Step Towards Reproducing OpenAI o1: Progress Report on the Steiner Open Source Models},
   url = {https://medium.com/@peakji/b9a756a00855},
   author = {Yichao Ji},
   month = {October},
   year = {2024}
}