北交開源o1代碼版！強(qiáng)化學(xué)習(xí)+蒙特卡洛樹搜索，源代碼、精選數(shù)據(jù)集以及衍生模型通通開源

西風(fēng) 2024-12-09 14:29:48 來源：量子位

西風(fēng) 發(fā)自 凹非寺

量子位 | 公眾號 QbitAI

北京交通大學(xué)研究團(tuán)隊(duì)悄默聲推出了一版o1，而且所有源代碼、精選數(shù)據(jù)集以及衍生模型都開源！

名為O1-CODER，專注于編碼任務(wù)。

團(tuán)隊(duì)認(rèn)為編碼是一個(gè)需要System-2思維方式的典型任務(wù)，涉及謹(jǐn)慎、邏輯、一步步的問題解決過程。

而他們的策略是將強(qiáng)化學(xué)習(xí)（RL）與蒙特卡洛樹搜索（MCTS）相結(jié)合，讓模型能夠不斷生成推理數(shù)據(jù)，提升其System-2能力。

實(shí)驗(yàn)中，團(tuán)隊(duì)有以下幾點(diǎn)關(guān)鍵發(fā)現(xiàn)：

• 當(dāng)推理正確時(shí)，基于偽代碼的推理顯著提升了代碼生成質(zhì)量
• 將監(jiān)督微調(diào)（SFT）與直接偏好優(yōu)化（DPO）相結(jié)合能夠提升測試用例生成效果
• 自我對弈強(qiáng)化學(xué)習(xí)為推理和代碼生成創(chuàng)造了持續(xù)改進(jìn)的循環(huán)機(jī)制

具體來說，團(tuán)隊(duì)采用了測試用例生成器，在經(jīng)過DPO后達(dá)到89.2%的通過率，相比初始微調(diào)后的80.8%有顯著提升；Qwen2.5-Coder-7B采用偽代碼方法實(shí)現(xiàn)了74.9%的平均采樣通過率，提升了25.6%。

網(wǎng)友直呼很需要這樣的模型。

O1-CODER，究竟長啥樣？

六步，逐步優(yōu)化o1

應(yīng)用于代碼生成的自我對弈強(qiáng)化學(xué)習(xí)面臨兩大挑戰(zhàn)：

• 結(jié)果評估，即如何評判生成代碼的質(zhì)量。與圍棋等任務(wù)不同，評估代碼需要在測試環(huán)境中運(yùn)行并驗(yàn)證。
• 定義思考和搜索行為，即確定過程獎(jiǎng)勵(lì)的對象和粒度。

對于第一個(gè)挑戰(zhàn)，團(tuán)隊(duì)提出訓(xùn)練一個(gè)測試用例生成器（TCG），根據(jù)問題和標(biāo)準(zhǔn)代碼自動(dòng)生成測試用例，為強(qiáng)化學(xué)習(xí)提供標(biāo)準(zhǔn)化的代碼測試環(huán)境和結(jié)果獎(jiǎng)勵(lì)。

對于第二個(gè)挑戰(zhàn)，他們采取”先思考后行動(dòng)“的方式：先通過詳細(xì)的偽代碼思考問題，再基于偽代碼生成最終的可執(zhí)行代碼。

這種方式的優(yōu)勢在于適應(yīng)性（同一偽代碼可對應(yīng)不同的具體實(shí)現(xiàn)）和可控粒度（通過調(diào)整偽代碼的細(xì)節(jié)程度控制推理/搜索行為的粒度）。

具體來說，研究團(tuán)隊(duì)提出了一個(gè)包含六個(gè)步驟的框架：

• 訓(xùn)練測試用例生成器（TCG），為代碼測試提供標(biāo)準(zhǔn)化的環(huán)境
• 利用MCTS生成包含推理過程的代碼數(shù)據(jù)
• 迭代微調(diào)策略模型，先生成偽代碼，再生成完整代碼
• 基于推理過程數(shù)據(jù)初始化過程獎(jiǎng)勵(lì)模型（PRM）
• 在TCG提供的結(jié)果獎(jiǎng)勵(lì)和PRM提供的過程獎(jiǎng)勵(lì)的雙重引導(dǎo)下，通過強(qiáng)化學(xué)習(xí)和MCTS更新策略模型
• 利用優(yōu)化后的策略模型生成新的推理數(shù)據(jù)，返回第4步迭代訓(xùn)練

兩階段訓(xùn)練測試用例生成器

在實(shí)驗(yàn)部分，研究人員詳細(xì)介紹了測試用例生成器的訓(xùn)練過程。

分為兩個(gè)階段：監(jiān)督微調(diào)（SFT）和直接偏好優(yōu)化（DPO）。

SFT階段的主要目標(biāo)是確保生成器的輸出符合預(yù)定義格式，以便準(zhǔn)確解析和提取生成的測試用例。訓(xùn)練數(shù)據(jù)來自TACO數(shù)據(jù)集。

DPO階段的目標(biāo)是引導(dǎo)模型生成符合特定偏好的測試用例，進(jìn)一步提高生成器的性能和可靠性。

這里采用了帶有人工構(gòu)建樣本對的DPO方法，構(gòu)建了一個(gè)偏好數(shù)據(jù)集。

實(shí)驗(yàn)表明，SFT階段過后，TCG在標(biāo)準(zhǔn)代碼上生成的測試用例通過率達(dá)到80.8%，DPO階段進(jìn)一步提升至89.2%，大幅改善了生成器產(chǎn)出可靠測試用例的能力。

偽代碼推理，引導(dǎo)模型進(jìn)行深度推理

特別值得一提的是，研究者引入了基于偽代碼的提示方法，將其作為引導(dǎo)模型進(jìn)行深度推理的“認(rèn)知工具”。

他們?yōu)榇硕x了三個(gè)關(guān)鍵行為：

• 使用偽代碼定義算法結(jié)構(gòu)：勾勒主要函數(shù)的結(jié)構(gòu)和接口，把握任務(wù)的整體框架
• 細(xì)化偽代碼：逐步明確每個(gè)函數(shù)的具體步驟、邏輯和操作
• 從偽代碼生成代碼：將偽代碼的結(jié)構(gòu)和邏輯精準(zhǔn)翻譯為可執(zhí)行代碼

在MBPP數(shù)據(jù)集上進(jìn)行的初步實(shí)驗(yàn)表明，盡管整體通過率（Pass@1）有所下降，但Average Sampling Pass Rate（ASPR）顯著提高。

表明結(jié)合偽代碼顯著改善了推理過程的質(zhì)量，特別是在細(xì)化通向正確輸出的路徑方面。這為后續(xù)的自監(jiān)督微調(diào)和強(qiáng)化學(xué)習(xí)提供了良好的起點(diǎn)。

自我對弈+強(qiáng)化學(xué)習(xí)

研究人員詳細(xì)描述了如何使用蒙特卡洛樹搜索（MCTS）來構(gòu)建步驟級別的過程獎(jiǎng)勵(lì)數(shù)據(jù)。

這個(gè)過程涉及到為每個(gè)問題形成一個(gè)推理路徑，該路徑由一系列推理步驟組成，并最終產(chǎn)生一個(gè)可執(zhí)行的代碼。在MCTS的路徑探索中，使用偽代碼提示策略來引導(dǎo)推理過程。當(dāng)達(dá)到終端節(jié)點(diǎn)時(shí)，就形成了一個(gè)完整的偽代碼推理路徑。

終端節(jié)點(diǎn)的獎(jiǎng)勵(lì)值是基于兩個(gè)關(guān)鍵指標(biāo)計(jì)算的：編譯成功率（compile）和測試用例通過率（pass）。

這些指標(biāo)被用來評估生成的代碼的質(zhì)量和正確性。

獎(jiǎng)勵(lì)值被反向傳播到路徑上的所有前序節(jié)點(diǎn)，為每個(gè)步驟分配一個(gè)獎(jiǎng)勵(lì)值。通過這種方式，構(gòu)建了推理過程數(shù)據(jù)集，為策略模型的初始化和訓(xùn)練提供了基礎(chǔ)。

過程獎(jiǎng)勵(lì)模型（PRM）的任務(wù)是為當(dāng)前步驟分配一個(gè)獎(jiǎng)勵(lì)值，以估計(jì)其對最終答案的貢獻(xiàn)。

在數(shù)據(jù)合成過程中使用的樹搜索方法可以組織成點(diǎn)式（point-wise）和成對式（pair-wise）兩種數(shù)據(jù)格式。

基于這些經(jīng)過驗(yàn)證的正確推理解，策略模型得到初始化。

接下來，過程獎(jiǎng)勵(lì)模型（PRM）開始發(fā)揮作用，評估每一步推理對最終答案的貢獻(xiàn)。在測試用例生成器（TCG）提供的結(jié)果獎(jiǎng)勵(lì)和PRM提供的過程獎(jiǎng)勵(lì)的雙重引導(dǎo)下，策略模型通過強(qiáng)化學(xué)習(xí)不斷改進(jìn)。

更新后的策略模型被用來生成新的推理數(shù)據(jù)，補(bǔ)充到現(xiàn)有數(shù)據(jù)集中，形成自我對弈的閉環(huán)。這個(gè)數(shù)據(jù)生成-獎(jiǎng)勵(lì)建模-策略優(yōu)化的迭代循環(huán)，確保了系統(tǒng)推理能力的持續(xù)提升。

論文鏈接：https://arxiv.org/pdf/2412.00154
參考鏈接：https://x.com/rohanpaul_ai/status/1864488583744377271?s=46&t=iTysI4vQLQqCNJjSmBODPw