Energy-Based Transformers are Scalable Learners and Thinkers

2025-07-02

AI Summary

Can systematic, System-2–like reasoning be achieved via purely unsupervised learning? This paper introduces Energy-Based Transformers (EBTs), which formulate prediction as an energy minimization problem: using a Transformer backbone, EBTs define a differentiable energy function between inputs and candidate outputs, enabling implicit reasoning via gradient-based optimization. Crucially, EBTs realize end-to-end “thinking” solely under unsupervised objectives—requiring no labels—while supporting cross-modal reasoning. Experiments demonstrate state-of-the-art scaling efficiency: up to 35% improvement across data, batch size, parameter count, FLOPs, and depth dimensions. On language understanding tasks, reasoning accuracy improves by 29%; in image denoising, EBTs surpass Diffusion Transformers with fewer forward passes. The core contribution lies in embedding systematic reasoning intrinsically within an unsupervised energy optimization framework.

Technology Category

• Machine Learning: Unsupervised & Self-Supervised Learning
• Reasoning under Uncertainty: Relational Probabilistic Models
• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

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.

Problem

• Generalize System 2 Thinking approaches for unsupervised learning
• Develop models verifying input-prediction compatibility without supervision
• Improve scalability and performance across modalities and tasks

Innovation

• Energy-Based Transformers for scalable learning
• Unsupervised learning with input-prediction compatibility
• Gradient descent-based energy minimization