一个提示、一个目标，以及用 Codex 运行 Colab！

几天前，@swyx 发布了一个挑战：用 JAX、Flax 和 Optax 编写一个大约 1000 万参数的 transformer，在免费的 Colab 中运行，并用它训练简单的加法。

我想看看 Codex 能否将截图变成一个我可以检查的 Colab 运行。

我几乎直接复制了提示，并添加了一个指令：/goal。在 Codex 中，goal 会将目标附加到线程上。

TL;DR： Codex 创建了一个 Colab 笔记本，包含一个 10,652,557 参数的仅解码器 transformer。它训练了 4000 步，耗时 1139.5 秒（约 19 分钟），在随机示例上达到了 99/100 的精确准确率，并在验证中达到了 199/200。

199/200 的验证结果来自监控的 Colab 输出。源笔记本还包括了 100 个示例的精确检查。

提示

提示要求提供模型代码，任务还需要浏览器操作：使用用户已认证的 Chrome 会话，创建新的 Colab 笔记本，选择运行时，监控长时间运行的单元格，检查保存的输出，并判断打印的结果是否支持该说法。

Colab 需要用户已登录的 Google 状态。Codex 使用了 Codex Chrome 扩展程序来实现这一点。该扩展程序允许 Codex 操作用户的 Chrome 浏览器，包括依赖 cookies、登录状态或仅浏览器 UI 的页面。

第一次失败

第一个问题是 Colab 编辑。笔记本编辑器有状态、自动保存、运行时对话框、键盘焦点和 Monaco 编辑器行为。

Codex 首先尝试直接重写单元格。有些编辑似乎生效了，然后又返回了过时的内容。实用的解决方法很简单：复制整个单元格，聚焦编辑器，全选，清除，粘贴，并验证活动单元格是否包含预期的标记。

在单元格运行时，浏览器检查也变的脆弱。Codex 使用保存的笔记本输出和 Drive 元数据作为更持久的证据链。

训练设置

词汇表正好是 13 个字符：数字 0-9、空格、+ 和 =。最长的可能示例是 999+999=1998，因此每个生成的序列都可以填充到固定长度 12。模型是一个仅解码器的 Flax transformer，包含六层、六个注意力头、宽度 384 和 MLP 宽度 1536。

初始化后，笔记本报告：

参数数量：10,652,557

这使得它在请求的 ~10M 范围内，而无需事后调整架构。

数据生成器从 0..999 中采样操作数，并格式化字符串，例如 123+456=579。它使用了均匀和长度平衡示例的混合，这有助于在训练期间足够频繁地出现较短的数字。

损失只针对答案：在看到 a+b= 后，模型学习生成答案数字，并在有空间时将尾随空格用作停止标记。

笔记本为 Codex 提供了足够的反馈，以判断训练是否有效。它打印了可见的 JAX 设备，生成了示例样本，运行了三个快速步骤，然后训练了 4000 个完整步骤，每 400 步打印损失、token 准确率和经过时间。

TPU，然后是 T4

会话中的一次较早运行似乎在免费的 TPU 运行时上完成了。在干净的监控重新运行期间，Colab 没有可用的 TPU 后端。

原始提示允许 T4 作为后备。Codex 切换了运行时，重新运行了笔记本，并监控训练直到完成。

在 T4 上重新运行确实有效：

JAX 设备：[CudaDevice(id=0)] 参数数量：10,652,557 第 4000 步 | 损失 0.0008 | token 准确率 100.0% | 1139.5s 随机示例精确准确率：99/100 (99.0%) 验证随机示例精确准确率：199/200 (99.5%)

token 准确率数字是训练期间同批次答案 token 的准确率。随机示例检查给出了序列级别的信号。

模型还通过了精选示例，这些示例涵盖了简单求和、进位、边界情况和最大长度结果：

0+0=0 7+5=12 12+34=46 99+1=100 123+456=579 999+999=1998 500+500=1000 908+95=1003

这些检查表明笔记本成功训练并泛化到了小的随机样本。要得出更有力的结论，需要更大的评估桶，或者对 0..999 + 0..999 的整个 1,000,000 个有序操作数对域进行穷举检查。

审计

运行后，我要求 Codex 使用子代理从不同角度审计结果：需求合规性、模型机制、数据生成、评估质量、Colab 可重现性、笔记本可用性以及怀疑接受度。

子代理是用于有界审查任务的独立代理运行。当主线程需要并行进行多个有针对性的检查，同时保留一个用于综合和决策的地方时，它们很有用。

审计发现模型设计对于此任务是合理的。因果掩码是正确的。移位后的语言模型目标与贪婪预测匹配。参数数量在范围内。词汇表和固定长度填充与提示匹配。

评估需要改进。最终的精确检查覆盖了 100 个随机示例和 200 个验证示例。两者都来自训练期间使用的同一组示例。训练 token 准确率来自当前批次。空格同时用作填充和停止标记。笔记本还需要固定的包版本、检查点以及更清晰的 CPU 快速失败路径。

P.S. 在真正有用之前还有很多改进空间，但这个方向是好的！

Codex 使用了什么

• /goal：将终点线附加到线程上。

• Chrome：赋予 Codex 访问已登录 Colab 会话的权限。

• 本地文件和 shell 检查：为 Codex 提供了一个稳定的源来与实时笔记本和保存的输出进行比较。

• 子代理：将最终审查拆分为有针对性的审计过程。

使其可靠

如果我再次运行，我会从五个方面使笔记本更可靠。

首先，使运行时明确。在顶部打印后端、加速器类型和包版本。在 CPU 上快速失败。将 TPU 不可用性作为清晰的分支处理，然后在可用时回退到 T4。

第二，将快速模式与完整模式分开。快速模式应运行几步并证明代码可运行。完整模式应运行实际的训练循环并以可预测的间隔打印进度。

第三，加强评估。保留小的随机检查以快速反馈，然后添加按操作数长度、进位行为、结果长度和边界情况（如 999+999）划分的桶。可选地添加对所有 1,000,000 个操作数对的穷举检查。

第四，保存紧凑的结果工件。包括配置、运行时、最终指标、精选示例、随机检查结果，理想情况下还包括检查点。笔记本输出很有用，但结果工件使后续比较更容易。

第五，使用 Codex 功能作为运行设置的一部分。使用明确的验收标准设置 /goal。使用 Chrome 进行经过身份验证的 Colab 工作。当浏览器变得脆弱时，要求 Codex 检查保存的输出。使用子代理进行怀疑性的最终审计。

对于未来的 Codex 运行：使用明确的验收标准设置目标，从快速测试开始，在真实环境中运行完整作业，保存收据，并要求进行怀疑性审计。

顺便说一下：这篇博文是 Codex 写的，它还创建了截图、推文以及周围的一切，我唯一做的就是格式化并点击发布！👀