Energy-Based Transformer sind skalierbare Lernende und Denker.

Zusammenfassung

Inferenzzeit-Berechnungstechniken, analog zum menschlichen System-2-Denken, sind in letzter Zeit populär geworden, um die Modellleistungen zu verbessern. Die meisten bestehenden Ansätze leiden jedoch unter mehreren Einschränkungen: Sie sind modalitätsspezifisch (z. B. funktionieren sie nur mit Text), problemspezifisch (z. B. verifizierbare Domänen wie Mathematik und Programmierung) oder erfordern zusätzliche Überwachung/Weiterbildung auf der Grundlage von unüberwachtem Vortraining (z. B. Verifizierer oder verifizierbare Belohnungen). In diesem Papier stellen wir die Frage: „Ist es möglich, diese System-2-Denkansätze zu verallgemeinern und Modelle zu entwickeln, die ausschließlich durch unüberwachtes Lernen lernen zu denken?“ Interessanterweise finden wir, dass die Antwort ja lautet, indem wir lernen, die Kompatibilität zwischen Eingaben und Kandidaten-Vorhersagen explizit zu verifizieren und dann Vorhersageprobleme als Optimierung in Bezug auf diesen Verifizierer neu zu formulieren. Konkret trainieren wir Energy-Based Transformers (EBTs) – eine neue Klasse von Energy-Based Models (EBMs) –, um jedem Eingabe- und Kandidaten-Vorhersage-Paar einen Energiewert zuzuweisen, wodurch Vorhersagen durch energiebasierte Gradientenabstiegsminimierung bis zur Konvergenz ermöglicht werden. Sowohl bei diskreten (Text) als auch bei kontinuierlichen (visuellen) Modalitäten stellen wir fest, dass EBTs während des Trainings schneller skalieren als der dominante Transformer++-Ansatz und eine bis zu 35 % höhere Skalierungsrate in Bezug auf Daten, Batch-Größe, Parameter, FLOPs und Tiefe erreichen. Während der Inferenz verbessern EBTs die Leistung mit System-2-Denken um 29 % mehr als der Transformer++ bei Sprachaufgaben, und EBTs übertreffen Diffusion Transformers bei der Bildentrauschung, während sie weniger Vorwärtsdurchläufe verwenden. Darüber hinaus stellen wir fest, dass EBTs bei den meisten nachgelagerten Aufgaben bessere Ergebnisse erzielen als bestehende Modelle bei gleicher oder schlechterer Vortrainingsleistung, was darauf hindeutet, dass EBTs besser generalisieren als bestehende Ansätze. Folglich sind EBTs ein vielversprechendes neues Paradigma für die Skalierung sowohl der Lern- als auch der Denkfähigkeiten von 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 Transformer sind skalierbare Lernende und Denker.

Energy-Based Transformers are Scalable Learners and Thinkers

Zusammenfassung

Support