概覽

O1 Nano 是一個開源項目，它實現(xiàn)了 OpenAI O1 模型系列的簡化版本。O1 模型是一個高級語言模型，它在訓(xùn)練和推理過程中整合了鏈?zhǔn)剿季S和強化學(xué)習(xí)。這個實現(xiàn)版本，稱為 O1-nano，專注于解決算術(shù)問題，以展示模型的能力。該項目受到了 Siraj Raval 在 YouTube 上的視頻的啟發(fā)。

主要特點

• 鏈?zhǔn)剿季S推理：模型生成完成令牌和內(nèi)部推理令牌，模擬思維過程。
• 強化學(xué)習(xí)：使用近端策略優(yōu)化（PPO）進(jìn)行訓(xùn)練。
• 多推理路徑：在生成過程中探索多條路徑，并選擇最佳路徑。
• 子任務(wù)生成：能夠?qū)?fù)雜問題分解為更小的子任務(wù)。
• 自適應(yīng)推理：包括在生成過程中修訂推理的機制。
• 大上下文窗口：支持高達(dá)128,000個令牌的上下文窗口。
• 內(nèi)部推理令牌：實現(xiàn)了用于推理的可丟棄內(nèi)部令牌。

文件

• train.py：包含模型架構(gòu)、訓(xùn)練循環(huán)和實用函數(shù)。
• test.py：提供了一個簡單的界面，用于與訓(xùn)練好的模型交互。

模型架構(gòu)

train.py 中的 O1Model 類定義了模型架構(gòu)：

• 嵌入層
• 位置編碼
• 多個變換器層
• 分別的完成和推理解碼器
• 用于強化學(xué)習(xí)的值頭
• 子任務(wù)生成頭

訓(xùn)練過程

訓(xùn)練過程結(jié)合了監(jiān)督學(xué)習(xí)和強化學(xué)習(xí)：

• 數(shù)據(jù)生成：算術(shù)問題即時生成。
• 監(jiān)督學(xué)習(xí)：模型被訓(xùn)練以預(yù)測正確的解決方案和推理步驟。
• 強化學(xué)習(xí)：使用 PPO 來優(yōu)化模型的策略，基于獎勵。
• 動態(tài)課程：根據(jù)訓(xùn)練進(jìn)度調(diào)整問題的難度。

使用方法

• 訓(xùn)練： 運行 python train.py 將訓(xùn)練模型500個周期，并將模型保存為 o1_model.pth。

• 測試： 運行 python test.py 將加載訓(xùn)練好的模型，并允許你輸入算術(shù)問題讓模型解決。

需求

• Python 3.7+
• PyTorch 1.8+

模型參數(shù)

• 嵌入維度：128
• 注意力頭數(shù)：8
• 變換器層數(shù)：4
• 丟棄率：0.1

詞匯表

模型使用了一個為算術(shù)操作定制的詞匯表，包括用于子任務(wù)和推理步驟的特殊令牌。

評估

模型根據(jù)其正確解決算術(shù)問題的能力進(jìn)行評估。評估指標(biāo)包括平均獎勵和處理的有效樣本數(shù)量。

• 規(guī)模：此實現(xiàn)比 OpenAI 描述的實際 O1 模型小。
• 任務(wù)多樣性：目前專注于算術(shù)；可以擴(kuò)展到更多樣化的任務(wù)。
• 自我修正：自我修正機制可以更加復(fù)雜。