单卡 5090 训练 1.5 小时:小型 Transformer 在 ARC-AGI-1 拿下 44%,推理成本仅 67 美分

单卡 5090 训练 1.5 小时:小型 Transformer 在 ARC-AGI-1 拿下 44%,推理成本仅 67 美分

班加罗尔的独立研究者 Mithil Vakde 发布了其 ARC-AGI 系列工作的第三篇博客:一个在单张 RTX 5090 上仅用 1.5 小时从零训练的小型 Transformer,在 ARC-AGI-1 公开评测集上取得 44% 的成绩,单次完整评测的推理成本只有 67 美分。同一模型在难度更高的 ARC-AGI-2 上拿到 7%,与专为此类谜题设计的 TRM、HRM 小模型基本持平,同时把大量参数规模大得多的通用大语言模型甩在身后。代码与权重照例开源,完整评测在一台消费级显卡上即可复现。

方法:测试时训练与序列建模

整体思路延续其上一代模型,但完成了一轮显著升级。每个输入-输出示例对被转换成 token 序列,由一个小型自回归 Transformer 进行训练;关键的测试时训练(test-time training)环节中,模型会对每道谜题在评测现场从零开始训练,训练数据包含该题的演示对与训练集谜题,评测集谜题的答案标签则始终隐藏。这种「先学通用变换规则、再在考场现场适应新规则」的范式,正是 ARC 基准希望考察的元学习能力。

为了让不同谜题之间的知识可以互相迁移,每道题被分配一个独立训练得到的加性嵌入向量;二维网格的位置信息则通过三维 RoPE 旋转位置编码注入序列模型。数据增强覆盖了颜色置换与二面角变换,让模型对格子着色与旋转翻转保持不变性。作者在博客中公开了全部技术细节与超参数,社区可以按图索骥地复现甚至改进这套流程。

为什么死磕样本效率

Vakde 在博客中直言,样本效率是他眼中当今 AI 最重要的问题,这项工作的目的有二:一是在只使用当代深度学习工具箱(即 Transformer)的前提下,摸清样本效率的极限在哪里;二是把迭代成本压到足够低,让任何背景的研究者都能快速试错。ARC 基准恰好是检验样本效率的理想靶场:公开评测集只有约一千道谜题,每道题使用不同的变换规则,训练集已经覆盖评测所需的全部基础概念,不需要外部先验。

这个项目的传播路径本身就值得一书。上一代模型的结果发布时,不少研究者认为数据量级下不可能达到该精度,相关讨论在 X 上迅速扩散,Lucas Beyer、Jeremy Howard、Rohan Anil 等知名研究者都下场参与了技术辨析,最终反而帮助作者确认了复现路径的可靠性。新一代结果发布后,讨论的焦点从「是否可能」转向了「测试时训练的上限究竟在哪」。

与 LLM 路线的对照

作者刻意把比较范围限定在同样执行测试时训练的方法之内,以避开「用海量合成数据或人类归纳偏置换分」的路线之争。主流 LLM 在 ARC 类基准上的成绩,往往建立在数十亿美元训练投入与巨量合成谜题预训练之上;而这条小模型路线把先验压缩到极限,用架构与训练策略去逼近人类在这类谜题上的表现。两条路线的对照,实质上是在回答「推理能力到底需要多少先验知识」这个根本问题。

成本维度的对比更加悬殊。前沿大模型跑一次完整 ARC 评测的推理开销通常以数十美元计,而这套方案只要 67 美分,训练投入是一张消费级显卡的 1.5 小时。作者特别强调低成本的开放意义:买不起算力集群的研究者、发展中国家的学生,都能在这个基准上做出有竞争力的工作。这对长期被算力门槛挡在门外的独立研究社区,是一个相当务实的信号。

接下来 Vakde 计划探索全新的研究想法,继续向样本效率的极限施压,同时维持低成本约束,让全世界的研究者都能参与迭代。无论这条路线最终能把 ARC-AGI-1 推到多高,它已经证明了一件事:在特定类型的抽象推理任务上,精心设计的小模型加测试时训练,可以用三个数量级更低的成本逼近乃至超越大模型的成绩。