迈向自动化AI研究的第一步 (12分钟阅读)

# 迈向自动化AI研究的第一步 - Recursive 文章来源：https://www.recursive.com/articles/first-steps-toward-automated-ai-research 今天，我们发布Recursive自动化AI研究系统的初步成果。在三个基准测试中，该系统取得了最先进的结果：固定预算的语言模型训练、小模型训练速度以及GPU内核优化。该系统针对目标指标自动执行研究循环：提出想法、实现想法、运行实验、验证结果，并利用所学知识选择下一个实验。它能在长时间跨度内并行运行多条研究线程，保留之前实验的有用上下文，合并有前景的分支，并在将性能提升视为真正进展之前，通过验证来排查奖励欺骗和方差问题。该系统设计为可扩展，并利用了开放式算法的原理，基于我们团队及其他团队在递归自我改进AI方面的前期工作。 我们选择的基准测试既具有实际重要性，又能提供紧密的反馈循环，着重考验AI进步的三大核心杠杆：更好的训练算法、更快的训练速度以及更高效的硬件利用。这些基准测试还非常适合自动化研究，因为它们具有明确的指标、相对较低的方差，并且评估器可以强化以抵御奖励欺骗。我们将在这些运行中产生的成果进行开源（https://github.com/recursive-org/first-steps-toward-automated-ai-research），以便他人检查并基于系统的输出进行构建。 | 基准测试 | 任务类型 | 指标 | 先前最先进水平 | Recursive改进 | | --- | --- | --- | --- | --- | | NanoChat Autoresearch | 在给定小计算预算下训练小语言模型至最高性能 | 验证BPB | 0.9372 | 0.9109 | 验证BPB降低0.0263，或达到相同损失的速度提升1.3倍 | | NanoGPT Speedrun | 尽可能快地将小语言模型训练到特定性能 | 达到3.28验证损失所需的训练时间 | 79.7秒 | 77.5秒 | 训练时间缩短2.2秒 | | SOL-ExecBench | 将GPU内核优化至硬件极限 | 235个内核的平均SOL分数 | 0.699 | 0.754 | 与最优性能估计1.0的差距缩小18% | ## 案例研究1：NanoChat Autoresearch Andrej Karpathy的NanoChat Autoresearch仓库（https://github.com/karpathy/autoresearch）是自动化研究系统的一个流行起点。任务是在单GPU上固定的五分钟预算内，将小语言模型训练到最低验证损失（以每字节比特数BPB衡量）。这对我们的系统来说是一个自然的测试，因为实验速度快、方差低，而且奖励欺骗相对容易检测。或许正是出于这些原因，一个公开的协作项目已经围绕这个设置形成。 Autoresearch@home（https://www.ensue-network.ai/lab/autoresearch）将原始设置扩展为协作环境，数十名人类和数百个他们的Agent共同改进性能。这为我们提供了一个比Karpathy单次通宵运行更强的比较点。我们想测试我们的系统能否改进整个由人类和Agent社区产生的解决方案。 我们的系统从与Autoresearch代码相同的初始种子解决方案开始。我们最初在NVIDIA H100 GPU上搜索，然后将发现的解决方案迁移到NVIDIA B200 GPU上运行，以便与公开结果公平比较。在移除先前最优autoresearch@home解决方案中的轻微奖励欺骗，并用10个随机种子评估后，其平均性能为0.9372 BPB。我们的系统找到了一个达到0.9109 BPB的解决方案，改进了0.0263 BPB。从另一个角度看，与最优autoresearch@home解决方案相比，我们的解决方案在大约1.3倍更少的训练时间内就达到了Karpathy原始通宵autoresearch BPB的质量。 **图1：NanoChat Autoresearch：随时间变化的验证BPB** 图1表明，Autoresearch从一个已经优化的模型开始，其中包含一些非平凡的设计决策。为此，我们测试了我们的系统是否也能从更弱的起点（一个朴素的初始实现，即带有AdamW的原始Transformer）做出改进。我们的系统将模型从1.059 BPB改进到0.9344 BPB（在NVIDIA B200 GPU上评估），再次优于autoresearch@home社区产生的最优解决方案。这不一定证明了独立的重新发现，因为底层模型可能知道许多公开技术，包括autoresearch@home社区使用或创建的技术，但这确实表明搜索过程能够从更弱的起点组装出一个有竞争力的训练栈。最终的解决方案也与公开的最优解决方案在多个方面有所不同。 **图2：NanoChat Autoresearch：按解决方案划分的最终验证BPB** **图3：NanoChat Autoresearch：随时间变化的训练损失** 我们的系统提出了哪些修改？最优解决方案并非由单一技巧驱动。它们结合了架构、短上下文记忆、辅助损失、注意力机制、优化器行为、权重衰减计划、编译器设置等方面的变化。最大的收益之一来自更丰富的短上下文记忆机制。基线已经使用了值嵌入；我们的系统通过哈希二元组和三元组嵌入表扩展了这个想法，并通过学习到的门控混入注意力值路径。这为模型提供了一种廉价的方式使用局部n-gram信息，而无需支付较慢的卷积或注意力密集型替代方案的时间成本。这与近期工作如DeepSeek Engram（https://arxiv.org/abs/2601.07372）相关，该工作探讨了哈希表作为稀疏性轴。在我们的设置中，哈希表可以为大约5000万参数的模型增加1-20亿个稀疏参数：大部分条目在任何给定批次中处于非活跃状态，且查找成本低廉。类似的哈希表和n-gram思想也出现在顶尖的NanoGPT Speedrun提交中。我们的系统通过将哈希二元组和三元组嵌入注入多层注意力的值向量中，并在不同层使用不同的哈希以减少重复碰撞，从而将该系列思想适应于固定预算设置。我们不知道先前有使用这种确切变体的工作。 下面的可展开框包含了系统解决方案中的精选技术细节。我们手动检查了这些输出，并借助AI辅助分析来理解技术并筛查奖励欺骗。我们可能仍然遗漏了内核优化中的错误（我们并非该领域的专家），但这正是关键所在：这里展示的想法来自系统，而非我们先前所擅长的领域。 ### 哈希表（https://www.recursive.com/articles/first-steps-toward-automated-ai-research#） 在起始解决方案的标准一元值嵌入之上，在我们的最优解决方案中，模型将每个二元组和三元组哈希到固定大小的表中，并通过学习到的、与输入相关的逐头门控，将查找出的向量混入注意力值路径——实际上是将经典的n-gram模型折叠到Transformer的值流中。 ```python for j, layer_i in enumerate(ve_layers): self.bigram_ves[str(layer_i)] = nn.ModuleList([ nn.Embedding(self.bigram_table_size, half_kv_dim), nn.Embedding(self.bigram_table_size, half_kv_dim), ]) self.bigram_hash_primes_per_layer[layer_i] = _decorr_bigram_primes[j] for j, layer_i in enumerate(sorted(self.trigram_ve_layers)): self.trigram_ves[str(layer_i)] = nn.ModuleList([ nn.Embedding(self.trigram_table_size, half_kv_dim), nn.Embedding(self.trigram_table_size, half_kv_dim), ]) self.trigram_hash_primes_per_layer[layer_i] = _decorr_trigram_primes[j] ``` ```python v = v + gate.unsqueeze(-1) * ve ## standard value embedding v = v + bg_gate.unsqueeze(-1) * bigram_ve ## additional bigram embedding from lookup table v = v + tg_gate.unsqueeze(-1) * trigram_ve ## additional trigram embedding from lookup table ``` 该解决方案还给不同的Transformer层分配了不同的哈希函数（使用互不相交的哈希质数对）。 ```python self.bigram_hash_primes_per_layer[layer_i] = _decorr_bigram_primes[j] self.trigram_hash_primes_per_layer[layer_i] = _decorr_trigram_primes[j] ``` 这意味着碰撞仍然会发生，但跨层以相同方式碰撞的可能性降低了。优化原始Transformer的运行时使用了与我们最优解决方案相同的一些技术，包括哈希表和平方ReLU MLP。但它也收敛到了一个不同（但同样有竞争力）的最终训练栈，包括token移位、评估前的权重平均以及字节级特征嵌入。这表明系统并非仅仅重复了另一次运行中的相同发现。下面的可展开框显示了原始Transformer运行中特有的一些修改。 ### 优化原始Transformer（https://www.recursive.com/articles/first-steps-toward-automated-ai-research#） 原始Transformer解决方案中的许多变化也出现在我们的最优解决方案中（来自我们的系统以Autoresearch初始种子代码启动），例如用Muon替换AdamW以及添加哈希表。其他一些改进并未在我们产生最优解决方案的主运行中出现，但却引起了我们的注意。第一个是因果token移位，它将前一个token的注意力投影Q和K混入当前token，每个维度附带一个学习到的系数。 ```python B, T, C = x.size() x_prev = F.pad(x[:, :-1, :], (0, 0, 1, 0)) q = self.c_q(x + self.q_shift_beta * x_prev).view(B, T, self.n_head, self.head_dim) k = self.c_k(x + self.k_shift_beta * x_prev).view(B, T, self.n_head, self.head_dim) v = self.c_v(x).view(B, T, self.n_head, self.head_dim) ``` 第二个是在token嵌入之后立即注入一组字节级特征。字节级特征表示关于token由哪些字节（例如单个字符）组成的信息。由相似字节组成的token将获得相似的字节级嵌入。字节特征嵌入矩阵的构建方式如下： ```python combined = torch.zeros(vocab_size, 769) for token_id in range(vocab_size): raw_bytes = tokenizer_enc.decode_single_token_bytes(token_id) # variable length if len(raw_bytes) > 0: for b in raw_bytes[:max_bytes]: # max_bytes=16 combined[token_id, b] += 1.0 / len(raw_bytes) # [0:256] byte-frequency histogram combined[token_id, 256 + raw_bytes[0]] = 1.0 # first byte one-hot combined[token_id, 512 + raw_bytes[-1]] = 1.0 # last byte one-hot combined[token_id, 768] = len(raw_bytes) / max_bytes # length feature torch.manual_seed(1337) proj = torch.randn(769, embed_dim) * 0.01 init_emb = combined @ proj # [vocab, embed_dim] self.embed = nn.Embedding(vocab_size, embed_dim) self.embed.weight.data.copy_(init_emb) # used only as the INIT ``` 这些嵌入随后在训练过程中通过梯度下降更新，并在token嵌入之后与二元组和三元组嵌入一起添加： ```python x_base = self.wte(idx) # token embedding gates = self.embed_mixer(x_base) # per-token gates over 4 sources [B,T,4] x = x_base x = x + gates[:,:,0:1] * bi_raw # bigram hash x = x + gates[:,:,1:2] * tri_raw # trigram hash x = x + gates[:,:,2:3] * self.byte_embed.get_raw(idx) # byte-content x = x + gates[:,:,3:4] * self.byte_boundary.get_raw(idx) # byte-boundary x = x + self.ssm_light(x) x = self.embed_ctx_norm(x) ``` 这些只是我们系统在这次运行中做出的部分改变。NanoChat展示了当要求我们的系统改进固定预算训练时，如何发现了许多相互叠加且对预算敏感的改进。下一个测试是，同样的过程是否能在经过多年的公开人工优化后，仍然在某个基准测试上找到提升空间。我们在NanoGPT Speedrun上进行了测试，其最佳公开解决方案已经过社区两年多的高度优化。 ## 案例研究2：NanoGPT Speedrun NanoGPT Speedrun（https://github.com/KellerJordan/modded-nanogpt）是一个类似的任务，但击败当前最先进水平要困难得多，因为一个庞大的社区已经对其解决方案进行了两年多的优化。该基准测试不问在固定时间预算内能达到多低的验证损失，而是问使用单个HGX H100 8-GPU节点，需要多快才能将一个小型GPT风格模型训练到在FineWeb文本数据集上达到固定的3.28验证损失。这是一项成熟的社区努力，截至目前已有83项人类创纪录的贡献提交到排行榜，以及数百个提议的PR。自2024年年中以来，训练时间从大约45分钟被压缩到79.7秒，这是一系列主要手工工程提交的结果。 鉴于当前解决方案已经高度优化，剩下的明显改进寥寥无几。从当前的领先解决方案出发，我们的系统找到了一组额外的优化，将训练时间从79.7秒减少到77.5秒，同时仍然满足排行榜的验证损失显著性要求（平均验证损失≤3.28，p<0.01）。¹（https://www.recursive.com/articles/first-steps-toward-automated-ai-research#footnote-1）这与近期的人工贡献相比，是类似或更大的改进。 ### 77.5秒解决方案：FP8注意力、探索噪声和谨慎嵌入（https://www.recursive.com/articles/first-steps-toward-automated-ai-research#） 以下变化是它创造并组合以达到77.5秒的创新示例（训练脚本和内核文件共约200行改动）。 **FP8注意力投影。** 人工记录运行将注意力QKV（query, key, value）和输出投影层保留在bf16（bfloat16，一种16位浮点格式），仅在LM头（预测下一个token的最终层）中使用FP8（8位浮点）。 系统将FP8推入了注意力路径（注意力层内部的计算）——一个自定义操作，在前向矩阵乘法中使用float8_e4m3（一种具有4位指数和3位尾数的FP8格式）以获得2倍张量核心吞吐量，同时将反向传播保留在bf16以保证稳定性，并且仅在训练期间使用（验证前向传播时恢复为bf16）： ```python class FP8LinearFunction(torch.autograd.Function): @staticmethod def forward(ctx, x, w): x_f8, w_f8 = x.to(torch.float8_e4m3fn), w.to(torch.float8_e4m3fn) y = torch._scaled_mm(x_f8, w_f8.T, out_dtype=torch.bfloat16, scale_a=ones, scale_b=ones, use_fast_accum=True) ctx.save_for_backward(x, w) return y @staticmethod def backward(ctx, grad_output): # backward stays in bf16 for precision x, w = ctx.saved_tensors return grad_output @ w, grad_output.T @ x # applied to the attention projections, training only: if self.training: qkv = fp8_linear_fn(x.view(-1, self.dim), qkv_w) else: q, k, v = F.linear(x, ...) # bf16 at eval ``` **优化器中的退火探索噪声（Langevin/SGLD）**（https://www.semanticscholar.org/paper/aeed631d6a84100b5e9a021ec1914095c66de415）。系统修改了NorMuon（https://arxiv.org/abs/2510.05491）更新，向正交化、方差减少的更新中注入零均值高斯噪声，噪声标准差按行缩放至更新本身的均方根（因此跨参数具有尺度不变性）。噪声在前50步中预热，然后在训练进行约四分之一时线性退火至零，因此最终解决方案倾向于干净地落在一个（希望更平坦的）盆地中： ```python def _add_exploration_noise(v_chunk, noise_t, red_dim): sigma = noise_t.to(v_chunk.dtype) # annealed scale (0-D CPU tensor) row_rms = v_chunk.float().square().mean(red_dim, keepdim=True).sqrt_() return v_chunk.add_(torch.randn_like(v_chunk) * (sigma * row_rms)) ``` **嵌入表上的谨慎（符号一致）Adam**。人工记录方案已经使用了谨慎权重衰减。系统将“谨慎”思想扩展到Adam更新步骤本身，针对二元组和val