I Transformer Basati sull'Energia sono Apprendenti e Pensatori Scalabili

Energy-Based Transformers are Scalable Learners and Thinkers

July 2, 2025

Autori: Alexi Gladstone, Ganesh Nanduru, Md Mofijul Islam, Peixuan Han, Hyeonjeong Ha, Aman Chadha, Yilun Du, Heng Ji, Jundong Li, Tariq Iqbal

cs.AI

Abstract

Le tecniche di calcolo al momento dell'inferenza, analoghe al pensiero di Sistema 2 umano, sono recentemente diventate popolari per migliorare le prestazioni dei modelli. Tuttavia, la maggior parte degli approcci esistenti presenta diverse limitazioni: sono specifici per una modalità (ad esempio, funzionano solo con il testo), specifici per un problema (ad esempio, domini verificabili come la matematica e la programmazione) o richiedono supervisione/addestramento aggiuntivo oltre al pre-addestramento non supervisionato (ad esempio, verificatori o ricompense verificabili). In questo articolo, ci chiediamo: "È possibile generalizzare questi approcci di pensiero di Sistema 2 e sviluppare modelli che imparano a pensare esclusivamente dall'apprendimento non supervisionato?" Interessantemente, scopriamo che la risposta è sì, imparando a verificare esplicitamente la compatibilità tra input e previsioni candidate, e poi riformulando i problemi di previsione come ottimizzazione rispetto a questo verificatore. Nello specifico, addestriamo Energy-Based Transformers (EBT) — una nuova classe di Energy-Based Models (EBM) — per assegnare un valore di energia a ogni coppia input-predizione candidata, consentendo previsioni attraverso la minimizzazione dell'energia basata sulla discesa del gradino fino alla convergenza. Sia nelle modalità discrete (testo) che continue (visive), scopriamo che gli EBT scalano più velocemente dell'approccio dominante Transformer++ durante l'addestramento, raggiungendo un tasso di scalabilità fino al 35% superiore rispetto a dati, dimensione del batch, parametri, FLOP e profondità. Durante l'inferenza, gli EBT migliorano le prestazioni con il pensiero di Sistema 2 del 29% in più rispetto al Transformer++ nei compiti linguistici, e gli EBT superano i Diffusion Transformers nella rimozione del rumore dalle immagini utilizzando meno passaggi in avanti. Inoltre, scopriamo che gli EBT ottengono risultati migliori rispetto ai modelli esistenti nella maggior parte dei task downstream a parità o peggiori prestazioni di pre-addestramento, suggerendo che gli EBT generalizzano meglio rispetto agli approcci esistenti. Di conseguenza, gli EBT rappresentano un nuovo paradigma promettente per scalare sia le capacità di apprendimento che di pensiero dei modelli.

English

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.