Energy-Based Transformers are Scalable Learners and Thinkers

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Gladstone, Alexi, Nanduru, Ganesh, Islam, Md Mofijul, Han, Peixuan, Ha, Hyeonjeong, Chadha, Aman, Du, Yilun, Ji, Heng, Li, Jundong, Iqbal, Tariq
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866908431505424384
author Gladstone, Alexi
Nanduru, Ganesh
Islam, Md Mofijul
Han, Peixuan
Ha, Hyeonjeong
Chadha, Aman
Du, Yilun
Ji, Heng
Li, Jundong
Iqbal, Tariq
author_facet Gladstone, Alexi
Nanduru, Ganesh
Islam, Md Mofijul
Han, Peixuan
Ha, Hyeonjeong
Chadha, Aman
Du, Yilun
Ji, Heng
Li, Jundong
Iqbal, Tariq
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2507_02092
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Energy-Based Transformers are Scalable Learners and Thinkers
Gladstone, Alexi
Nanduru, Ganesh
Islam, Md Mofijul
Han, Peixuan
Ha, Hyeonjeong
Chadha, Aman
Du, Yilun
Ji, Heng
Li, Jundong
Iqbal, Tariq
Machine Learning
Artificial Intelligence
Computation and Language
Computer Vision and Pattern Recognition
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.
title Energy-Based Transformers are Scalable Learners and Thinkers
topic Machine Learning
Artificial Intelligence
Computation and Language
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2507.02092