o1不是唯一路径!MIT新研究:测试时训练,模型推理能力大幅提升

  克雷西发自凹非寺

  量子位公众号 QbitAI

  o1 不是通向大模型推理的唯一路径!

  MIT 的新研究发现,在测试时对大模型进行训练,可以让推理水平大幅提升。

  在挑战超难的 ARC 任务时,准确率最高可提升至原来的 5.83 倍。

  这样的表现不仅优于 GPT-4 和 Claude,如果与其他推理方法相结合,还能超越人类的平均水准。

  OpenAI o1 团队成员Noam Brown表示,o1 的大规模计算可能不是最好的方法,很高兴看到有学者在提高推理能力上探索新的方法。

  在测试中训练模型

  不同于传统的先训练后测试模式,测试时训练(Test-Time Training,TTT)在部署阶段面对新的测试样本时,不直接用训练好的模型去推理

  在推理之前,测试样本自身携带的信息,会通过快速的训练过程被用于调整模型参数。

  总体来说,TTT 过程中一共有三个关键阶段——训练数据生成、模型适应范式设计以及推理阶段的策略

  数据生成的核心是将测试任务中蕴含的输入输出对关系,通过数据增强的方式最大限度地利用,可具体分为两个步骤。

  首先是基于 leave-one-out 构造新的任务。

  对于包含K个输入输出对的测试任务,依次将每个样本留出作为测试样本,其余K-1 个作为训练样本,由此构造出K个新的 TTT 训练任务。

  这样就可以从一个测试任务出发,构造出K个结构一致但内容互补的新任务,从而扩充了 TTT 训练数据。

  在此基础上,作者还进行了数据增强,主要包括对输入输出施加各类几何变换,以及打乱训练样本对的顺序。

  经过这一步,TTT 训练集的规模可以得到显著扩大。

  整个 TTT 数据构造过程可高度自动化,不依赖人工标注。

  利用构造好的 TTT 数据集,就可以对预训练好的语言模型进行测试时训练。

  考虑到测试时的资源限制,作者采用了参数高效的 LoRA,为每个测试任务学习一组独立的 adapter 参数,附加在预训练模型的每一层之上,通过一个低秩矩阵与原始权重相乘起到调节作用。

  过程中还额外加入了对所有前缀序列的预测,目的是通过在各种长度的演示样本上都计算损失,鼓励模型尽早地从少量信息中总结出抽象规律,从而提高鲁棒性。

  最后,为了实现 TTT 效果的最大化,作者在推理阶段应用了数据增强和集成学习策略。

  推理过程中,先利用一系列预定义的几何变换算子(如旋转、翻转等)扩充原始输入,生成若干等价视角下的输入变体。

  之后将每个变体输入并行地送入 LoRA-tuned 模型,独立完成预测,然后再对齐和还原到原始输入空间,由此得到一组成对的预测。

  在成对预测的基础上,通过分两层投票的方式完成集成融合:

  • 第一层在每种变换内部进行投票,选出置信度最高的 Top-3 个预测;
  • 第二层在不同变换的 Top-3 预测之间进行全局投票,选出最终的 Top-2 作为输出。

  这一推理策略,既通过数据增强引入了输入的多样性,又用分层投票的方式对不同来源的预测进行了结构化的组合,进一步提升了 TTT 方法的效果。

  ARC 任务准确率最高升至 6 倍

  为了评估 TTT 方法的效果,研究团队以 8B 参数的 GPT-3 作为基础模型进行了测试。

  如果不使用 TTT 仅进行微调,模型在 ARC 数据集上的准确率只有 18.3%,加入 TTT 后提升到 47.1%,增长率达到了 157%。

  另外,作者还从 ARC 数据集中随机选择了 80 个任务作为子集进行了测试。

  测试发现,TTT 方法对于 1B 模型的提升效果更加明显,调整后模型的准确率接近调整前的 6 倍。

  并且在调整前后,1B 和 8B 两个规模的模型之间的相对差距也在缩小。

  进一步地,作者还将 TTT 方法与之前在 ARC 任务上取得优异成绩的BARC(Bootstrapping Approach for Reward model Construction)方法进行了比较和结合。

  具体来说,作者首先独立运行这两个系统,得到它们在每个测试任务上的输出。

  如果两者输出完全一致,则直接认为推理结果是正确的;

  如果输出不一致,则看 BARC 是否能够生成确定的、唯一覆盖所有测试样本的解题程序,若是则认为 BARC 的输出更可靠;

  反之,如果 BARC 生成了多个候选程序但无法确定最优解,或者干脆无法生成任何满足约束的程序,则认为 TTT 的输出更可靠。

  两种方式配合使用后,取得了 61.9% 的 SOTA 成绩,已经超过了人类的平均水平

  One More Thing

  根据作者在推文中的介绍,在这篇论文发布前,一个叫做 MindsAI 的团队已经发现使用了相同的技术。

  利用 TTT 技术,该团队已经用 58% 的正确率取得了 ARC 挑战的第一名。

  作者的论文发布之后,MindsAI 团队领导者 Jack Cole 也发文进行了祝贺:

很高兴,我们掀起了这场对 TTT 的兴趣风暴。

  同时,Jack 还推荐了另一名研究 TTT 的学者——斯坦福大学华人博士后Yu Sun,表示他的研究值得被关注。

  Sun 的个人主页显示,他针对测试时训练进行了大量研究,相关成果入选过 ICML、NeurIPS、ICLR 等多个顶级会议。

  论文地址:

  https://ekinakyurek.github.io/papers/ttt.pdf