Modelli Cognitivamente Ispirati Basati sull'Energia per la Rappresentazione del Mondo

Cognitively Inspired Energy-Based World Models

June 13, 2024

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

cs.AI

Abstract

Uno dei metodi predominanti per addestrare modelli del mondo è la previsione autoregressiva nello spazio di output dell'elemento successivo di una sequenza. Nel Natural Language Processing (NLP), ciò si traduce in Large Language Models (LLM) che prevedono il token successivo; in Computer Vision (CV), ciò assume la forma di modelli autoregressivi che prevedono il frame/token/pixel successivo. Tuttavia, questo approccio differisce dalla cognizione umana sotto diversi aspetti. Innanzitutto, le previsioni umane sul futuro influenzano attivamente i processi cognitivi interni. In secondo luogo, gli esseri umani valutano naturalmente la plausibilità delle previsioni riguardanti stati futuri. Basandosi su questa capacità e, in terzo luogo, valutando quando le previsioni sono sufficienti, gli esseri umani allocano una quantità dinamica di tempo per fare una previsione. Questo processo adattivo è analogo al pensiero di Sistema 2 in psicologia. Tutte queste capacità sono fondamentali per il successo degli esseri umani nel ragionamento e nella pianificazione di alto livello. Pertanto, per affrontare le limitazioni dei tradizionali modelli autoregressivi che mancano di queste capacità simili a quelle umane, introduciamo i Modelli del Mondo Basati sull'Energia (EBWM). L'EBWM prevede l'addestramento di un Modello Basato sull'Energia (EBM) per prevedere la compatibilità di un determinato contesto e di uno stato futuro previsto. In questo modo, l'EBWM consente ai modelli di raggiungere tutti e tre gli aspetti della cognizione umana descritti. Inoltre, abbiamo sviluppato una variante del tradizionale trasformatore autoregressivo adattata per i modelli basati sull'energia, denominata Trasformatore Basato sull'Energia (EBT). I nostri risultati dimostrano che l'EBWM scala meglio con i dati e le ore di GPU rispetto ai tradizionali trasformatori autoregressivi in CV e che l'EBWM offre promettenti risultati iniziali di scalabilità in NLP. Di conseguenza, questo approccio offre un percorso entusiasmante verso l'addestramento di futuri modelli capaci di pensiero di Sistema 2 e di ricerca intelligente attraverso spazi di stati.

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.