Energy-Based Transformers are Scalable Learners and Thinkers

Paper · arXiv 2507.02092 · Published July 2, 2025
Novel LLM Architectures

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 (unnormalized probability) value to every input and candidate-prediction pair, enabling predictions through gradient descent-based energy minimization until convergence. This formulation enables System 2 Thinking to emerge from unsupervised learning, making it modality and problem agnostic. 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 (i.e., extra computation) 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 System 2 Thinking with EBTs yields larger performance improvements on data that is farther out-of-distribution, and 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.

Synthesis notes that discuss concepts related to this paper 2

Notes in this collection most semantically related to the paper, ranked by cosine similarity.

Lines of inquiry this paper opens 24

Research framings built by reading the notes related to this paper — the questions it feeds into.

Do language models learn genuine linguistic structure or surface-level patterns? What determines when reasoning models help or hurt task performance? How do neural networks internally organize representations and reasoning? How can iterative computation overcome fundamental limitations in model capability? What causes policy entropy collapse during reasoning reinforcement learning? Does recurrent computation enable new reasoning capabilities fixed-depth architectures cannot achieve? How should inference-time compute be allocated for different task difficulties? How do representation structure properties emerge and change in neural networks? Why does task-specific specialization harm generalization and future learning? How should training curricula adapt difficulty as models improve? Does reasoning learned in one domain transfer to other fields? Can activation steering control reasoning without retraining models? How does explicit structure affect the depth and consistency of language model reasoning? What graph topologies best support reasoning and model generalization? Do base models contain latent reasoning that training unlocks? Can AI reasoning scale in latent space without tokens? How does AI assistance during learning affect knowledge retention and transfer? Does chain-of-thought reasoning represent genuine model computation?