我花了几个小时逐行查看 /karpathy/autoresearch 仓库。 “人工智能代理进行研究”的角度吸引了所有的关注，但我认为更有趣的是训练脚本内部的内容以及使搜索循环紧凑的工程决策。这是我读过的最密集的单文件训练设置之一。 让我从使整个项目成为可能的事情开始：时间预算固定为 300 秒的实际时间。不是固定的步骤，不是固定的标记，不是固定的浮点运算。是实际时间的秒数。这听起来像是一个小细节，但这正是自主循环能够工作的原因。代理可以将模型扩大 3 倍，将批量大小减半，替换为完全不同的架构，结果仍然可以直接与其他实验进行比较，因为它们都在同一 GPU 上进行了恰好 5 分钟的训练。如果你固定步骤，那么更大的模型每秒获得的梯度更新会更少，你会不公平地惩罚它。如果你固定标记，你会遇到同样的问题。固定墙时间意味着你在问正确的问题：在这个硬件和这么多时间的情况下，你能产生的最佳模型是什么？其他一切都是自由变量。代理可以探索模型大小与吞吐量与收敛速度之间的完整帕累托曲面，而不会被评估协议混淆。 指标也经过精心选择。它是每字节位数，而不是交叉熵损失。交叉熵依赖于你的词汇大小。一个有 32k 标记的模型和一个有 8k 标记的模型，即使它们以同样的方式压缩数据，损失值也会有很大不同。bpb 通过在 nats 中求和每个标记的交叉熵，求和目标标记的 utf-8 字节长度，并将 nats 每字节转换为每字节位数，从而消除了这一点。因此，即使代理更改了影响有效标记分布的内容，比较仍然是公平的。这两个选择，固定墙时间和不变的词汇指标，将本来会是一个混乱的不可比较的搜索转变为一个干净的优化问题。 现在是模型本身。它是一个 GPT，但有许多现代技巧值得理解。首先，处处使用 RMSnorm。在块输入（预归一化）上，以及在注意力点积之前的查询和键上。这个 QK-norm 事情很重要，因为没有它，q 和 k 的范数在训练过程中可能会无限增长，导致注意力 logits 变得尖锐，softmax 饱和。对 q 和 k 进行归一化可以保持点积在一个稳定的范围内，无论网络有多深或训练动态如何演变。注意力本身是 FA 3，通过内核库加载。它使用 varunneal 在 hopper (sm_90) 上的实现，并在较旧的 GPU 上回退到社区构建。注意力模式是“SSSL”，这意味着三层滑动窗口注意力（窗口 = 序列长度的一半），后面跟着一层完整的因果注意力，重复。这是你在 mistral 和 gemma2 中看到的稀疏到密集的模式。 局部注意力层在计算上是便宜的，因为注意力矩阵是带状的，周期性全局层允许信息在整个上下文中流动。通过 8 层和 4 字符模式，你得到层 0,1,2 局部，层 3 全局，层 4,5,6 局部，层 7 全局。最后一层无论模式如何都被强制为全局。 值嵌入的事情是微妙的，我认为被低估了。每一层都有自己的嵌入表，完全独立于主标记嵌入，直接将标记 ID 映射到值维度向量。这些通过学习的门混合到注意力值中：v = v + 2 * sigmoid(W_gate @ x:32) * ve。门权重是零初始化的，因此 sigmoid(0) = 0.5，乘以 2 得到 1.0，这是一个中性的起点。在训练过程中，模型可以根据隐藏状态的前 32 个维度学习放大或抑制每个头的值嵌入。这来自 ResFormer 的工作线，直觉是它给注意力提供了一个直接的快捷方式来获取标记身份。值向量可以携带关于“这个位置是什么标记”的信息，而不必让这些信息在早期层的残差流转换中存活。它本质上是一个从输入直接到注意力值的跳过连接，门控使模型可以决定何时有用。 每层的残差流上还有可学习的标量：x = lambda_residi * x + lambda_x0i * x0，其中 x0 是来自层 0 的归一化嵌入。每一层可以独立控制它对运行残差与原始输入的倾听程度。残差 lambda 从 1.0 开始，x0 lambda 从 0.1 开始。这是“解耦残差”思想的软版本。在标准变压器中，残差流是所有先前层输出的总和，随着你深入，它变得越来越污染。让每一层访问干净的原始嵌入意味着它不必学习“撤销”早期层以恢复低级信息。logits 通过 tanh(logits/15)*15 软限制在 15 以内，这防止了模型在训练早期过于自信，因为表示仍然是嘈杂的。 但老实说，整个文件中最有趣的部分是优化器。MuonAdamW 是一个组合优化器，根据参数组分派不同的更新规则。嵌入（标记嵌入、值嵌入、去嵌入头）和每层标量获得标准的 AdamW，每组有不同的学习率。差异很大。嵌入 lr 是 0.6，去嵌入 lr 是 0.004，这差了 150 倍，而且这是故意的。嵌入矩阵看到每一个标记，需要积极更新。去嵌入矩阵是最终表示上的线性探测，受益于稳定性。嵌入、值嵌入和去嵌入的学习率都按 (d_model / 768)^(-0.5) 进行缩放，这是一个受 muP 启发的修正。随着模型宽度的变化，这些学习率会调整，以保持特征学习动态的尺度不变。每层 lambda 的标量学习率单独处理，不会进行这种缩放。 变压器中的 2D 权重矩阵、注意力投影和 mlp 权重，使用 Muon，这里变得真正有趣。muon 获取梯度，应用 nesterov 动量，然后运行 newton-schulz 迭代以近似梯度矩阵的极分解。极分解将矩阵 G 分解为 G = U * S，其中 U 是正交的，S 是对称正半定的。muon 计算 U，即最接近梯度的正交矩阵，并将其用作更新方向。newton-schulz 迭代进行 5 步。对于高矩阵（行数多于列数），A = X^T @ X 然后 X -> aX + X @ (bA + cA^2)。对于宽矩阵，A = X @ X^T 然后 X -> aX + (bA + cA^2) @ X。系数是从预计算中硬编码的。他们称之为“极地快车”。整个过程通过 torch.compile 编译为一个单一的融合内核。 这有什么重要性？因为对于权重矩阵，弗罗贝纽斯范数梯度（adam 和 sgd 使用的）在几何上是错误的。权重矩阵的“正确”最陡下降方向是最小化损失的方向，前提是更新具有单位谱范数，而不是单位弗罗贝纽斯范数。正交极因子正好给你这个。在实践中，这意味着 muon 进行的有效更新要大得多，因为它没有在缩放奇异值上浪费步长。它只旋转它们。这就是为什么 muon 在变压器权重矩阵上收敛速度显著快于 adam。muon 确实维护每个元素的动量缓冲区（与参数形状相同，堆叠在每个形状组中），但与 adam 不同，它不跟踪每个元素的二次矩。二次矩估计是在正交化后按行或按列进行的，而不是按元素进行的。这就是 NorMuon 的作用。 在基础 muon 之上还有 NorMuon，一种方差减少方案。在正交化后，它计算按行（或按列，具体取决于纵横比）的二次矩估计，维护其指数移动平均，并重新缩放更新，以便每个输出维度获得自己的自适应步长。这本质上是将 adam 自适应性理念应用于正交化坐标系，而不是原始参数空间。权重衰减也不是标准的。它是“谨慎的”，意味着它只衰减 muon 更新方向与参数符号一致的参数：mask = (g * params) >= 0。这避免了已知的失败模式，即权重衰减将参数推向零，而与更新的愿望相悖，这可能会使训练不稳定。 我欣赏的一个小细节是：在第一次训练步骤之后，代码调用 gc.collect()、gc.freeze()、gc.disable()，以完全关闭 Python 的垃圾收集器。Python 的 GC 定期运行并导致 ~500ms 的停顿。当你的总预算是 300 秒，每一步可能是 300ms 时，随机的 GC 暂停几乎会让你损失 2 个训练步骤。他们每 5000 步手动触发 gc.collect() 作为折衷。这是你通过分析真实训练运行并注意到神秘的吞吐量下降时才能学到的事情。 前 11 步（0 到 10）也不计入时间预算。这是热身阶段，torch.compile 完成其工作，CUDA 内核被 JIT 编译。如果没有这个排除，不同实验将根据特定模型配置的编译时间长短获得不同数量的“真实”训练。再次强调，这似乎是一个小设计选择，但对于使实验可比较至关重要。 现在放大一下。实际的自研究循环是：代理读取 program.md（描述其工作的 markdown 文件），修改 train.py，提交，运行 5 分钟，检查 val_bpb 是否改善，保留或恢复，重复。program.md 明确表示“永远不要停止”。代理无限期运行，直到人类终止它。每小时约 12 次实验，晚上睡觉时约 100 次。 ...