Modèles cognitifs du monde basés sur l'énergie inspirés par la cognition

Cognitively Inspired Energy-Based World Models

June 13, 2024

papers.authors: Alexi Gladstone, Ganesh Nanduru, Md Mofijul Islam, Aman Chadha, Jundong Li, Tariq Iqbal

cs.AI

papers.abstract

L'une des méthodes prédominantes pour entraîner des modèles du monde est la prédiction autorégressive dans l'espace de sortie de l'élément suivant d'une séquence. En traitement du langage naturel (NLP), cela se traduit par des modèles de langage de grande taille (LLMs) prédisant le prochain token ; en vision par ordinateur (CV), cela prend la forme de modèles autorégressifs prédisant la prochaine image/token/pixel. Cependant, cette approche diffère de la cognition humaine à plusieurs égards. Premièrement, les prédictions humaines sur l'avenir influencent activement les processus cognitifs internes. Deuxièmement, les humains évaluent naturellement la plausibilité des prédictions concernant les états futurs. Sur la base de cette capacité, et troisièmement, en évaluant quand les prédictions sont suffisantes, les humains allouent une quantité de temps dynamique pour faire une prédiction. Ce processus adaptatif est analogue à la pensée de type Système 2 en psychologie. Toutes ces capacités sont fondamentales pour le succès des humains dans le raisonnement et la planification de haut niveau. Par conséquent, pour répondre aux limitations des modèles autorégressifs traditionnels qui manquent de ces capacités humaines, nous introduisons les Modèles du Monde Basés sur l'Énergie (EBWM). EBWM consiste à entraîner un Modèle Basé sur l'Énergie (EBM) pour prédire la compatibilité d'un contexte donné et d'un état futur prédit. Ce faisant, EBWM permet aux modèles d'atteindre les trois aspects de la cognition humaine décrits. De plus, nous avons développé une variante du transformateur autorégressif traditionnel adapté aux modèles basés sur l'énergie, appelé le Transformateur Basé sur l'Énergie (EBT). Nos résultats démontrent que EBWM s'adapte mieux aux données et aux heures GPU que les transformateurs autorégressifs traditionnels en CV, et que EBWM offre un début prometteur de mise à l'échelle en NLP. Par conséquent, cette approche offre une voie passionnante pour entraîner des modèles futurs capables de pensée de type Système 2 et de recherche intelligente à travers les espaces d'états.

English

One of the predominant methods for training world models is autoregressive prediction in the output space of the next element of a sequence. In Natural Language Processing (NLP), this takes the form of Large Language Models (LLMs) predicting the next token; in Computer Vision (CV), this takes the form of autoregressive models predicting the next frame/token/pixel. However, this approach differs from human cognition in several respects. First, human predictions about the future actively influence internal cognitive processes. Second, humans naturally evaluate the plausibility of predictions regarding future states. Based on this capability, and third, by assessing when predictions are sufficient, humans allocate a dynamic amount of time to make a prediction. This adaptive process is analogous to System 2 thinking in psychology. All these capabilities are fundamental to the success of humans at high-level reasoning and planning. Therefore, to address the limitations of traditional autoregressive models lacking these human-like capabilities, we introduce Energy-Based World Models (EBWM). EBWM involves training an Energy-Based Model (EBM) to predict the compatibility of a given context and a predicted future state. In doing so, EBWM enables models to achieve all three facets of human cognition described. Moreover, we developed a variant of the traditional autoregressive transformer tailored for Energy-Based models, termed the Energy-Based Transformer (EBT). Our results demonstrate that EBWM scales better with data and GPU Hours than traditional autoregressive transformers in CV, and that EBWM offers promising early scaling in NLP. Consequently, this approach offers an exciting path toward training future models capable of System 2 thinking and intelligently searching across state spaces.