Modelos Cognitivos del Mundo Basados en Energía Inspirados en la Cognición

Cognitively Inspired Energy-Based World Models

June 13, 2024

Autores: Alexi Gladstone, Ganesh Nanduru, Md Mofijul Islam, Aman Chadha, Jundong Li, Tariq Iqbal

cs.AI

Resumen

Uno de los métodos predominantes para entrenar modelos del mundo es la predicción autorregresiva en el espacio de salida del siguiente elemento de una secuencia. En el Procesamiento del Lenguaje Natural (PLN), esto toma la forma de Modelos de Lenguaje de Gran Escala (LLMs) que predicen el siguiente token; en Visión por Computadora (CV), esto se manifiesta como modelos autorregresivos que predicen el siguiente fotograma/token/píxel. Sin embargo, este enfoque difiere de la cognición humana en varios aspectos. Primero, las predicciones humanas sobre el futuro influyen activamente en los procesos cognitivos internos. Segundo, los humanos evalúan naturalmente la plausibilidad de las predicciones sobre estados futuros. Basándose en esta capacidad, y tercero, al determinar cuándo las predicciones son suficientes, los humanos asignan una cantidad dinámica de tiempo para hacer una predicción. Este proceso adaptativo es análogo al pensamiento de Sistema 2 en psicología. Todas estas capacidades son fundamentales para el éxito de los humanos en el razonamiento y la planificación de alto nivel. Por lo tanto, para abordar las limitaciones de los modelos autorregresivos tradicionales que carecen de estas capacidades similares a las humanas, introducimos los Modelos del Mundo Basados en Energía (EBWM). EBWM implica entrenar un Modelo Basado en Energía (EBM) para predecir la compatibilidad de un contexto dado y un estado futuro predicho. Al hacerlo, EBWM permite que los modelos logren los tres aspectos de la cognición humana descritos. Además, desarrollamos una variante del transformador autorregresivo tradicional adaptado para modelos basados en energía, denominado Transformador Basado en Energía (EBT). Nuestros resultados demuestran que EBWM escala mejor con los datos y las horas de GPU que los transformadores autorregresivos tradicionales en CV, y que EBWM ofrece un prometedor escalamiento temprano en PLN. En consecuencia, este enfoque ofrece un camino emocionante hacia el entrenamiento de futuros modelos capaces de pensamiento de Sistema 2 y búsqueda inteligente a través de espacios de estados.

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.