Energy-Based Transformers zijn schaalbare leerders en denkers.

Samenvatting

Inferentie-tijd rekentechnieken, analoog aan menselijk Systeem 2 Denken, zijn recentelijk populair geworden voor het verbeteren van modelprestaties. De meeste bestaande benaderingen hebben echter verschillende beperkingen: ze zijn modaal-specifiek (bijvoorbeeld alleen werkend in tekst), probleem-specifiek (bijvoorbeeld verifieerbare domeinen zoals wiskunde en codering), of vereisen extra supervisie/training bovenop ongesuperviseerde voorafgaande training (bijvoorbeeld verificators of verifieerbare beloningen). In dit artikel stellen we de vraag: "Is het mogelijk om deze Systeem 2 Denken-benaderingen te generaliseren, en modellen te ontwikkelen die leren denken uitsluitend vanuit ongesuperviseerd leren?" Interessant genoeg vinden we dat het antwoord ja is, door expliciet te leren de compatibiliteit tussen invoer en kandidaat-voorspellingen te verifiëren, en vervolgens voorspellingsproblemen te herformuleren als optimalisatie met betrekking tot deze verificator. Specifiek trainen we Energy-Based Transformers (EBTs) – een nieuwe klasse van Energy-Based Models (EBMs) – om een energie-waarde toe te kennen aan elk invoer- en kandidaat-voorspellingspaar, waardoor voorspellingen mogelijk worden via energie-minimalisatie op basis van gradient descent tot convergentie. Over zowel discrete (tekst) als continue (visuele) modaliteiten heen, vinden we dat EBTs sneller schalen dan de dominante Transformer++-benadering tijdens training, met een tot 35% hogere schaalsnelheid ten opzichte van data, batchgrootte, parameters, FLOPs en diepte. Tijdens inferentie verbeteren EBTs de prestaties met Systeem 2 Denken met 29% meer dan de Transformer++ bij taal-taken, en presteren EBTs beter dan Diffusion Transformers bij beeldruisverwijdering terwijl ze minder forward passes gebruiken. Verder vinden we dat EBTs betere resultaten behalen dan bestaande modellen bij de meeste downstream-taken gegeven dezelfde of slechtere voorafgaande trainingsprestaties, wat suggereert dat EBTs beter generaliseren dan bestaande benaderingen. Bijgevolg zijn EBTs een veelbelovend nieuw paradigma voor het schalen van zowel de leer- als denkcapaciteiten van modellen.

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.

Energy-Based Transformers zijn schaalbare leerders en denkers.

Energy-Based Transformers are Scalable Learners and Thinkers

Samenvatting

Support