Les Transformeurs Basés sur l'Énergie sont des Apprenants et Penseurs Évolutifs

Résumé

Les techniques de calcul au moment de l'inférence, analogues à la pensée de type Système 2 chez l'humain, ont récemment gagné en popularité pour améliorer les performances des modèles. Cependant, la plupart des approches existantes souffrent de plusieurs limitations : elles sont spécifiques à une modalité (par exemple, fonctionnant uniquement pour le texte), spécifiques à un problème (par exemple, des domaines vérifiables comme les mathématiques et la programmation), ou nécessitent une supervision/formation supplémentaire en plus d'un pré-entraînement non supervisé (par exemple, des vérificateurs ou des récompenses vérifiables). Dans cet article, nous posons la question suivante : « Est-il possible de généraliser ces approches de pensée de type Système 2 et de développer des modèles qui apprennent à penser uniquement à partir d'un apprentissage non supervisé ? » Fait intéressant, nous constatons que la réponse est oui, en apprenant à vérifier explicitement la compatibilité entre les entrées et les prédictions candidates, puis en reformulant les problèmes de prédiction comme une optimisation par rapport à ce vérificateur. Plus précisément, nous entraînons des Transformers à Base d'Énergie (EBTs) — une nouvelle classe de modèles à base d'énergie (EBMs) — pour attribuer une valeur d'énergie à chaque paire d'entrée et de prédiction candidate, permettant ainsi des prédictions via une minimisation de l'énergie basée sur la descente de gradient jusqu'à convergence. À la fois pour les modalités discrètes (texte) et continues (visuelles), nous constatons que les EBTs s'adaptent plus rapidement que l'approche dominante Transformer++ pendant l'entraînement, atteignant un taux de mise à l'échelle jusqu'à 35 % supérieur en termes de données, taille de lot, paramètres, FLOPs et profondeur. Pendant l'inférence, les EBTs améliorent les performances avec la pensée de type Système 2 de 29 % par rapport au Transformer++ sur les tâches linguistiques, et les EBTs surpassent les Transformers de Diffusion pour le débruitage d'images tout en utilisant moins de passes avant. De plus, nous constatons que les EBTs obtiennent de meilleurs résultats que les modèles existants sur la plupart des tâches en aval, même avec un pré-entraînement identique ou inférieur, suggérant que les EBTs généralisent mieux que les approches existantes. Par conséquent, les EBTs représentent un nouveau paradigme prometteur pour améliorer à la fois les capacités d'apprentissage et de réflexion des modèles.

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.

Les Transformeurs Basés sur l'Énergie sont des Apprenants et Penseurs Évolutifs

Energy-Based Transformers are Scalable Learners and Thinkers

Résumé

Support