基于能量的变压器是可扩展的学习者和思考者

推理时计算技术类似于人类的系统2思维（System 2 Thinking），近年来在提高模型性能方面变得越来越受欢迎。然而，大多数现有的方法存在几个局限性：它们通常是特定模态的（例如，仅适用于文本），特定问题的（例如，可验证领域如数学和编程），或者需要在无监督预训练基础上进行额外的监督/训练（例如，验证器或可验证奖励）。本文中，我们提出了一个问题：“是否可以将这些系统2思维方法泛化，并开发出仅通过无监督学习就能学会思考的模型？”有趣的是，我们发现答案是肯定的，通过显式验证输入与候选预测之间的兼容性，并将预测问题重新表述为针对该验证器的优化问题。具体而言，我们训练了一种新的能量基模型（Energy-Based Models, EBMs）——能量基变压器（Energy-Based Transformers, EBTs），为每一对输入和候选预测分配一个能量值，通过基于梯度下降的能量最小化直至收敛来实现预测。在离散模态（文本）和连续模态（视觉）上，我们发现EBTs在训练过程中比主流的Transformer++方法扩展得更快，在数据量、批大小、参数、浮点运算次数和深度方面实现了高达35%更高的扩展率。在推理阶段，EBTs在语言任务上的表现比Transformer++提高了29%，并且在图像去噪任务上优于扩散变压器（Diffusion Transformers），同时使用的前向传播次数更少。此外，我们发现，在预训练性能相同或较差的情况下，EBTs在大多数下游任务上取得了更好的结果，这表明EBTs比现有方法具有更好的泛化能力。因此，EBTs是一种有前景的新范式，可以同时扩展模型的学习能力和思考能力。