O1 Nano：OpenAI O1模型系列的简化开源版本

概览

O1 Nano 是一个开源项目，它实现了 OpenAI O1 模型系列的简化版本。O1 模型是一个高级语言模型，它在训练和推理过程中整合了链式思维和强化学习。这个实现版本，称为 O1-nano，专注于解决算术问题，以展示模型的能力。该项目受到了 Siraj Raval 在 YouTube 上的视频的启发。

主要特点

• 链式思维推理：模型生成完成令牌和内部推理令牌，模拟思维过程。
• 强化学习：使用近端策略优化（PPO）进行训练。
• 多推理路径：在生成过程中探索多条路径，并选择最佳路径。
• 子任务生成：能够将复杂问题分解为更小的子任务。
• 自适应推理：包括在生成过程中修订推理的机制。
• 大上下文窗口：支持高达128,000个令牌的上下文窗口。
• 内部推理令牌：实现了用于推理的可丢弃内部令牌。

文件

• train.py：包含模型架构、训练循环和实用函数。
• test.py：提供了一个简单的界面，用于与训练好的模型交互。

模型架构

train.py 中的 O1Model 类定义了模型架构：

• 嵌入层
• 位置编码
• 多个变换器层
• 分别的完成和推理解码器
• 用于强化学习的值头
• 子任务生成头

训练过程

训练过程结合了监督学习和强化学习：

• 数据生成：算术问题即时生成。
• 监督学习：模型被训练以预测正确的解决方案和推理步骤。
• 强化学习：使用 PPO 来优化模型的策略，基于奖励。
• 动态课程：根据训练进度调整问题的难度。

使用方法

• 训练： 运行 python train.py 将训练模型500个周期，并将模型保存为 o1_model.pth。

• 测试： 运行 python test.py 将加载训练好的模型，并允许你输入算术问题让模型解决。

需求

• Python 3.7+
• PyTorch 1.8+

模型参数

• 嵌入维度：128
• 注意力头数：8
• 变换器层数：4
• 丢弃率：0.1

词汇表

模型使用了一个为算术操作定制的词汇表，包括用于子任务和推理步骤的特殊令牌。

评估

模型根据其正确解决算术问题的能力进行评估。评估指标包括平均奖励和处理的有效样本数量。

• 规模：此实现比 OpenAI 描述的实际 O1 模型小。
• 任务多样性：目前专注于算术；可以扩展到更多样化的任务。
• 自我修正：自我修正机制可以更加复杂。