Reed's News
← 返回精选

1.5小时训练的小型Transformer击败众多LLM

AI 72 porridgeraisin 2026/9/1 3216 字 原文 ↗

我用RTX 5090从零训练了一个小型Transformer,耗时仅1.5小时

它性能超越多数大语言模型(LLM),得分与TRM/HRM持平

这是我初代模型的升级版本

更快、更强、成本更低,且依旧开源

在ARC-2上也取得了7%的成绩

Twitter讨论,GitHub代码

ARC-1公开评估结果

这是ARC-AGI系列研究的第3篇博客。前两篇:[博客2](https://mvakde.github.io/(https:/mvakde.github.io/blog/new-pareto-frontier-arc-agi/#implementation-details)、[博客1](https://mvakde.github.io/blog/why-all-ARC-solvers-fail-today/)。

很多人认为我上次的结果不可能实现。研究成果受到顶尖学者关注,在X平台迅速传播,例如Lucas Beyer、Jeremy Howard、Rohan Anil等人都参与了讨论,还有众多研究者留下评论。

为什么做这项研究?

我认为样本效率是当前AI领域最重要的待解问题,而我想要攻克它。

本研究的目标有二:(1)探索Transformer及现有深度学习方法下,样本效率的极限在哪里;(2)降低研发成本,让迭代更快、更便宜。

ARC是测试样本效率的绝佳基准,原因如下:

  • 样本量极少(仅1000道谜题),但任务空间维度极高
  • 属于元学习基准,每道谜题遵循不同规则,但包含共通概念
  • 所需先验知识极少:测试集中的所有概念,训练集均已覆盖
  • 人类极易上手,即使是缺乏GPU资源的AI研究者也能参与
  • 基准仍未被充分挖掘(针对数据效率而言,需排除使用大量合成数据或人类归纳偏置的LLM及相关方法)

接下来,我将探索新的研究思路,突破现有极限。我会尽可能控制成本,让全球任何研究者都能参与其中。

技术细节

模型原理

整体思路与上次类似(完整技术细节见此),但新增了多项改进。核心方法概述如下:

  • 将每个输入输出对转换为token序列,由小型Transformer进行自回归训练。测试阶段,会在训练集和测试集谜题(测试标签隐藏)上从零开始训练。
  • 为实现跨任务学习,给每个谜题分配独立的可学习附加嵌入。由于每个序列包含两个二维网格,采用3D RoPE(旋转位置编码)嵌入来学习位置信息。
  • 对序列进行颜色和二面体变换增强。推理时,先对测试输入做增强处理,再对输出做逆增强变换,最终提交出现次数最多的两个结果(AAIVR策略)。

相较于初代模型的改进

本次核心目标是优化模型架构与算法,提升样本效率。

得分提升的关键因素:

  • 采用现代架构(用SwiGlu替代GELU,用RMSNorm替代LayerNorm等)
  • 增加数据多样性,优化数据洗牌策略
  • 模型扩容:从4层增至8层

成本降低的关键因素:

  • 大幅减少数据增强操作(样本效率显著提升!)
  • 优化器从AdamW替换为Normuon
  • 训练阶段使用变长FlashAttention,推理阶段使用FlexAttention核

一项重要变化是不再对输入token进行训练,损失函数仅包含输出token(即转为监督式训练)。这一改动让得分从40%小幅提升至44%,但背后原因我尚未完全理清,可能与模型有限的容量有关。

我还通过加入ARC-2中不重叠的任务扩充了训练数据,且严格避免了数据泄露。若移除新增数据,模型仍能达到约40%的得分,但所需算力会翻倍。

背景说明:ARC-2包含773道ARC-1重复谜题和347道全新谜题。由于ARC-1的多数测试谜题在ARC-2中重复出现,若直接用ARC-2训练会导致数据泄露,甚至能得到100%的得分。我通过仔细过滤掉773道重复谜题,彻底避免了这一问题!

此外还有诸多小改动,分别在性能或速度上带来了增量提升。完整改动列表见此。

有趣的模型表现

由于不再训练输入token,当前方法属于监督式训练。奇怪的是,测试损失反而上升了,但得分却更高!同时模型表现更稳定,得分波动更小。

如今很多研究者通过追求小数据集上的最低验证损失来提升样本效率,这本身无可厚非,但本研究揭示了这种思路的一种失效模式。

我认为无监督式训练在某些场景下仍会更优,目前正在开展相关评估。

在使用NorMuon之前,我尝试过基础版Muon。它的训练速度确实比AdamW快得多,但训练后期损失(及得分)会陷入停滞,无法收敛。我发现此时大幅调低动量或学习率能有所改善,但我不想依赖这类手动调整。换成NorMuon后,这个问题就消失了。

变量 ablation 实验

对性能贡献最大的因素似乎是优质的表征方式(3D RoPE+任务专属嵌入)。

RoPE与任务专属嵌入的ablation实验

  • 训练输入token的表现略差,得分约39%
  • 仅用ARC-1+ConceptARC训练,得分与原模型相近,约40%
  • 将3D RoPE改为1D RoPE,得分骤降至约24%
  • 移除任务专属嵌入,得分骤降至约24%
  • 采用CompressARC模式(每个任务单独从零训练,无监督),得分降至约18%
  • 采用CompressARC模式但改为监督式训练,得分约15%

其他ablation实验的最佳得分

如何参与贡献?

代码已开源,欢迎大家修改优化,提升得分或降低成本。(请勿增加训练数据)

尝试冲击65%的得分——无需大量修改即可实现。依据如下:我汇总多次运行中解决的所有任务,已达到55%的得分,还有不少任务处于"接近解决"的状态。提供几个思路:

  • RoPE会混合位置信息与内容信息,可能影响性能。PoPE(位置感知旋转编码)表现应该相当或更优,也可以尝试发明新的位置嵌入方式
  • 模型架构显然还有进一步现代化的空间

通过手写GPU代码,成本有望降低10倍,也可通过架构优化实现。

最后,想办法彻底移除数据增强操作。(我并不想用这种方法,也不认同那些认为它合理的观点)已有几种可行思路,但难点在于要同时保持低训练成本。

杂谈

说实话,我没想到仅用Transformer就能达到45%的得分,原以为需要全新的思路。更没想到能以如此低的成本和算力实现。Ablation实验显示,即使去掉数据增强或合成数据,模型仍能保留相当一部分性能。现在我确信,在Transformer框架内就能达到65%的得分。

我不理解为什么其他人没发现这一点。这不过是采用了最直观表征方式的Transformer而已。这个基准已经开放6年,关注度高,还有百万美元奖金!也许研究者低估了深度学习的能力?也许实验成本过高,无法充分开展ablation实验?或是被LLM吸引了注意力,或是过度依赖现有工具?

附录

知名学者对我研究方法的批评与验证

上次的研究成果在X平台引发热议,众多资深学者展开了正反两方的辩论,包括Jeremy、Lucas、Susan、Andreas、Yoav等人的讨论线程。我在此汇总所有批评及我的回应。

在测试谜题上训练属于作弊/“测试集训练”

  • 这种说法错误。“测试集训练”特指在测试数据的标签上训练,而本研究从未使用测试标签。
  • 此外,ARC是元学习基准,本来就应该从测试谜题中学习。
  • 术语说明:ARC包含训练谜题集和测试谜题集,每个谜题包含示例对和测试对,每一对由输入网格+输出网格组成。
  • 在ARC中,“标签”特指测试谜题中测试对的输出网格。
  • 这些标签完全未被用于训练,全程处于隐藏状态。你甚至可以提前删除它们,不影响实验。

在测试谜题的输入上训练会泄露信息

  • 这种说法错误。这种方法称为转导推理,从Vapnik时代起就已被研究。
  • 再者,在持续学习的研究背景下,“忽略测试输入”这一教条并不合理。
  • 其他元学习方法通常是先训练元学习算法,再通过思维链(CoT)或循环修改隐变量来部署学习。而我的方法是直接修改单个前向函数的权重来进行元学习,本质上并无不同。
  • 补充:在本次44%得分的新模型中,我已移除了输入训练,因为它的表现略差。

即使允许在测试谜题的输入上训练,也应禁止使用测试输入本身

  • 这种说法错误。同样适用上述“转导推理”的论据。
  • 元学习基准的转导性体现在两个层面:
    • 从训练谜题到测试谜题的转导
    • 在单个测试谜题内部,从示例对到测试对的转导
  • 这一批评恰好可以用第二个层面的转导性来回应。

这违反测试规则

  • 这种说法错误。
  • 规则原文是“测试者不得提前知晓测试内容”,人们误以为这禁止了测试时训练(TTT),但实际上规则约束的是设计AI系统的人类,而非AI系统本身——目的是防止人类基于测试集设计归纳偏置。
  • 对于ARC社区的活跃研究者而言,测试时训练被允许且鼓励是共识。Steven和Chew的评论也澄清了这一点及其他相关疑问。
  • 测试时训练符合元学习基准的核心精神,完全合规!

你没有计算训练成本

  • 这种说法错误。我统计的是全生命周期算力,包括从初始化开始的模型训练成本,加上在所有任务上运行推理的总成本。总成本确实仅为67美分,可以查看vast.ai平台上RTX 5090运行2小时的价格。

传统的测试时训练是逐个任务进行的,同时在所有测试任务上训练不现实

  • 这一批评有一定道理,但需要具体分析。
  • 我认同现实中很少遇到所有问题同时出现的场景,即使有(比如考试),人类通常也只能一次解决一个问题。
  • 但不能因为人类不具备某种能力,就否定AI拥有这种能力的合理性。否则我们也可以说LLM不现实,因为人类无法在整个互联网上训练,也无法以如此快的速度读取token。
  • 此外,人类是否真的只能一次处理一个任务也尚不明确——人类或许可以同时从多种感官渠道接收不同数据并学习,就像本模型这样。

按任务分摊成本,会因为同时训练所有测试任务而摊薄成本,与其他模型对比不公平

  • 这个观点合理。我的回应如下:
    • 赛事组织者就是这样对比所有模型的,包括同样同时训练所有测试任务的TRM。
    • 我已经更为严谨:我统计了训练和推理的全部成本,而LLM及其他模型通常会排除预训练/离线训练成本。
  • 我现在已调整了对比方式:(a)展示全生命周期算力成本而非单任务成本;(b)仅与TRM、HRM和CompressARC对比,不与LLM及其他方法对比;(c)新增了训练方式可比的ablation实验。

回应关于ARC-AGI本身的批评

上次发布研究后,很多人对ARC-AGI本身提出了质疑。其中有些观点合理,但更多是Chollet早已多次解答过的问题:

  • ARC到底测试什么?(流体智力)
  • 我们为什么要关注ARC?(流体智力尚未被完全解决)
  • 解决ARC-AGI不等于实现AGI(没人这么说过)
  • ARC一直在改变规则/是针对LLM对抗性设计的(两种说法都不对)

Chollet的论文及相关推文是很好的参考资料。总结他的观点:该基准旨在测试流体智力,他认为这是AGI的必要非充分条件。解决ARC-1/2意味着模型具备一定的流体智力,但这并非上限。基准的作用不是宣告AGI实现,而是指明值得研究的方向。规则从未改变:ARC-1早于LLM出现,ARC-2在ChatGPT发布前就已公布,ARC-3则在ARC-2被充分研究前就已宣布。他对ARC上的进展感到欣慰,因为这记录了AI的发展。

我关注ARC主要是因为它可用于测试样本效率——这是当前AI领域重要的未解决问题!它还是一个设计精良的元学习基准,GPU资源有限的研究者也能参与。从历史来看,它能很好地揭示LLM的优势与不足。另外,尽管ARC关注度高、奖金丰厚,却在6年内都未被充分攻克,如今我们知道深度学习能在ARC-1/2上取得极佳表现,这一点也很有意思。

我认为有几个批评是合理的:

  • 应禁止在ARC-1/2中使用合成数据

    • 这违背了基准的初衷,但当前多数高分模型都依赖大量合成数据
    • 合成数据降低了解决谜题所需的流体智力门槛
    • 这种做法在2024年之前是合理的,因为当时深度学习得分很低
    • 现在我们知道,只要有足够训练数据,LLM/深度学习能学习任何任务
  • 这会让ARC成为测试样本效率的绝佳基准,与ARC-3形成互补

  • 如何禁止合成数据?方法很简单:

    • 禁止离线训练/预训练,模型必须在提交后从零开始训练
    • 之前认为这不可能,所以没有相关规则,但我的模型证明这是可行的
    • 这样就能保证没有合成数据被使用
    • 让不同模型的对比更公平,否则像LLM这类模型可以通过海量离线训练“刷分”(由于基准存在时间较长,已出现很多类ARC数据集)
  • 用单一排行榜对比不同类型模型不合理,会引发三个问题(解决方案:分开展示图表):

    • X轴是单任务成本,但仅统计在线算力成本。部分模型(如LLM)有大量未被统计的离线预训练成本。它们可以用无限的训练算力让测试集数据分布与训练集一致,因此应单独评估。
    • 对于像我的模型、TRM和HRM这类同时训练所有测试任务的模型,按任务分摊成本毫无意义。
    • 在公开评估集上对比LLM得分毫无意义,因为公开谜题的答案已在互联网上可获取,且可能已被用于训练。
  • 赛事组织者从TRM和HRM的结果中过早得出结论,将成功归因于循环机制+深度监督。我认为这种偏见源于他们假设纯深度学习无法解决ARC(例如基础LLM在ARC-2上表现依旧糟糕),但我不同意这一点。

  • 测试规则的表述可以更清晰,避免歧义。(详见此处)

我认为其他方法存在的误区

认为循环机制是下一个突破点(例如HRM、TRM、Arcprize博客)

我理解循环机制的吸引力,但目前没有足够的ablation实验证明其必要性。我的模型表明,无需循环机制也能达到相同性能。已证实的循环机制唯一优势是可以在不增加内存占用的前提下提升算力。

HRM/TRM的误导性宣传:我不认同TRM自称是7M参数模型,因为它实际上训练了约1亿以上的嵌入权重。这属于误导,让模型更像一个查找表,也让人质疑其性能的来源。最稳妥的说法应该是7M“活跃”参数。HRM也存在同样问题,而且两者都未提及这一点!

基于LLM的ARC方法不再展现新能力:

LLM在ARC排行榜上的攀升曾带来很多启发(如下文所述),但我认为如今从它们的ARC-1/ARC-2得分中已学不到太多东西:

  • LLM得分的提升现在主要由训练后优化驱动(下一节有相关证据),且可能与合成数据的数量直接相关。它们是在学习如何解决ARC任务,而非学习通用抽象推理。
  • 新得分中存在太多干扰因素,无法得出有效结论。总体而言,基于基准对比LLM属于不严谨的科研行为(例如不同模型针对基准使用的训练数据量不同)。
  • 对于LLM而言,只有私有评估集的得分才有意义。它们在公开排行榜上的得分毫无价值,因为公开谜题的答案已在互联网上可获取,且可能已被用于训练。
  • 在LLM之上使用工具库(harness)提升性能的做法对我意义不大。所有训练后优化的核心技术都掌握在前沿实验室手中,他们自己就能开发工具库。我认为持续学习不太可能通过工具库解决。

违背“苦涩教训”的取巧手段

我之前就论证过,合成数据和数据增强是不可取的,在模型中设计归纳偏置同样不可取。我们必须依赖这类取巧手段才能提升基准性能,说明仍有突破性的方法等待发现。我希望更多人尝试减少这类违背“苦涩教训”的技巧。

从LLM在ARC-AGI上的表现得到的启示

LLM如今已在该基准的第1版和第2版上达到饱和。从它们的进展中我得出以下结论:

ARC-AGI早于LLM出现。最初LLM在该基准上表现极差,说明预训练无法赋予通用推理能力,且LLM在人类极易完成的任务上可能表现糟糕。

OpenAI的O1模型取得75%的得分是LLM的一大突破,表明只要有足够数据,LLM可以在训练后阶段学习任何任务。我认为这就是Sholto Douglas经常主张的观点。

ARC-2发布后,所有LLM(包括具备推理能力的模型)的进展都被重置。这说明即使是训练后优化也无法赋予通用推理能力,否则在ARC-1上表现良好的模型应该能自动在ARC-2上取得好成绩。

(本质上,这些模型是在学习如何解决ARC谜题,而非学习通用抽象推理,它们的任务得分取决于该任务在训练数据中的代表性。另外,我甚至不确定“通用推理”是否真的存在?也许人类的能力也是高度专业化的)

此后,具备推理能力的LLM在ARC-2上稳步提升。人们常认为这意味着模型的通用推理能力增强,但他们忽略了基础模型的得分仍停留在个位数。结合其他证据,这表明:

  • ARC-2的得分由训练后优化驱动(可能很大程度上取决于合成ARC数据的数量)
  • 实验室在“刷基准分”(可能因为客户关注基准性能)
  • LLM在样本效率方面毫无优势

别误会,我非常看好LLM。ARC-2的趋势表明,性能会随着算力和数据的增加持续提升。推理成本的大幅降低也令人惊叹。

完整改动列表

改变训练动态的改动

  • 优化器从仅使用AdamW改为NorMuon+辅助AdamW
  • 学习率调度从预热+余弦退火改为WSD调度(预热百分比,保持,然后线性衰减至下限)
  • 将LayerNorm替换为RMSNorm
  • 前馈网络(FFN)从Linear -> GELU -> Linear改为SwiGLU风格的门控FFN(分块+SiLU门)
  • 权重衰减从“仅针对非注意力线性层”改为显式分组衰减:注意力权重、token嵌入、任务/二面体嵌入各有独立的权重衰减控制参数
  • 训练目标从outputs["loss"](无监督式输入+输出语言模型损失)改为仅使用outputs["output_loss"](监督式损失)
  • 训练批处理从基于长度的智能分桶改为完全随机批处理(推理路径保留分桶策略)
  • 训练中丢弃不完整的批次(强制设置drop_last=True)
  • 数据集构建支持/使用更广泛的数据源(ARC-1、ARC-2、ConceptARC、可选的跨数据集过滤任务、提交/私有模式),改变了训练数据的构成
  • 颜色增强从全局 epoch 级别的单一排列改为每个样本的增强组合(颜色+二面体变换)
  • 颜色排列范围改为排除仅出现在输出中的颜色,而非盲目排列1-9号颜色
  • 增强生成通过哈希对任务内的变换输入去重(提升了唯一样本的多样性)
  • 增强选择改为按 epoch 循环,候选顺序随机打乱(每个循环内不重复选择)
  • 超参数调整:优化器超参数、训练轮数、增强上限/类型、模型深度(n_layers)、数据集路径
  • 新增dihedral_embedding,并将其加入token条件(仅带来非常轻微的性能提升)

不改变训练动态的速度优化

  • 训练批次从填充后的[B,S]格式改为使用cu_seqlens的打包token流(训练路径无填充token)
  • 注意力路径从填充SDPA掩码改为支持打包变长的FlashAttention(cu_seqlens),并在推理时使用FlexAttention核
  • 二面体增强从离线数据集扩展改为在批处理整合时在线选择增强方式
  • 构建时的训练拆分从("train","test")改为("train",)

其他改动

  • 断点续训行为优化:优化器切换/超参数变化检测现在可以重置/重新预热调度,改变了续训的动态过程
  • 调度器步进改为基于训练时的分数epoch进度,而非纯步数的余弦进度

TODO:添加引用