Abstract
Inference-time computation techniques, analogous to human System 2 Thinking, have recently become popular for improving model performances. However, most existing approaches suffer from several limitations: they are modality-specific (e.g., working only in text), problem-specific (e.g., verifiable domains like math and coding), or require additional supervision/training on top of unsupervised pretraining (e.g., verifiers or verifiable rewards). In this paper, we ask the question "Is it possible to generalize these System 2 Thinking approaches, and develop models that learn to think solely from unsupervised learning?" Interestingly, we find the answer is yes, by learning to explicitly verify the compatibility between inputs and candidate-predictions, and then re-framing prediction problems as optimization with respect to this verifier. Specifically, we train Energy-Based Transformers (EBTs) -- a new class of Energy-Based Models (EBMs) -- to assign an energy value to every input and candidate-prediction pair, enabling predictions through gradient descent-based energy minimization until convergence. Across both discrete (text) and continuous (visual) modalities, we find EBTs scale faster than the dominant Transformer++ approach during training, achieving an up to 35% higher scaling rate with respect to data, batch size, parameters, FLOPs, and depth. During inference, EBTs improve performance with System 2 Thinking by 29% more than the Transformer++ on language tasks, and EBTs outperform Diffusion Transformers on image denoising while using fewer forward passes. Further, we find that EBTs achieve better results than existing models on most downstream tasks given the same or worse pretraining performance, suggesting that EBTs generalize better than existing approaches. Consequently, EBTs are a promising new paradigm for scaling both the learning and thinking capabilities of models.
Abstract (translated)
类似于人类的系统2思维，推理时的计算技术最近因其在提升模型性能方面的潜力而变得流行。然而，大多数现有的方法存在若干局限性：它们只适用于特定模式（例如仅限于文本），针对特定问题领域（如数学和编程验证域），或者需要在无监督预训练基础上额外的监督或训练（例如校验器或可验证奖励）。在这篇论文中，我们提出了一个问题：“是否有可能将这些系统2思维方法进行泛化，并开发出仅仅通过无监督学习就能学会思考的模型？”令人惊讶的是，我们发现答案是肯定的。具体而言，我们的研究发现通过学习明确地验证输入与候选预测之间的兼容性，并重新框架为基于校验器的优化问题，可以实现这一目标。也就是说，我们训练了一种新型的能量基础模型（EBMs），即能量基础变换器（EBTs）来给每个输入和候选预测对分配一个能量值，使通过基于梯度下降的最小化直至收敛来进行预测成为可能。 在离散（文本）和连续（视觉）模式上，我们发现与目前主导的方法Transformer++相比，在训练过程中，EBTs的学习速率更快。具体来说，在数据量、批次大小、参数数量、FLOPs以及模型深度等方面，EBTs的扩展速度比Transformer++快35%以上。在推理阶段，相对于Transformer++方法，EBTs在语言任务上的系统2思维性能提高了29%，并且在图像去噪上优于扩散变换器，同时使用了更少的前向传递次数。此外，我们还发现，在给定相同或较差预训练表现的情况下，EBTs在大多数下游任务中的表现比现有的模型更好，这表明EBTs具有更好的泛化能力。 因此，能量基础变换器（EBTs）为增强模型的学习和推理能力提供了一种有前景的新范式。
URL
https://arxiv.org/abs/2507.02092
