Cognitief Geïnspireerde Energie-Gebaseerde Wereldmodellen

Cognitively Inspired Energy-Based World Models

June 13, 2024

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

cs.AI

Samenvatting

Een van de overheersende methoden voor het trainen van wereldmodellen is autoregressieve voorspelling in de uitvoerruimte van het volgende element in een reeks. In Natural Language Processing (NLP) neemt dit de vorm aan van Large Language Models (LLMs) die het volgende token voorspellen; in Computer Vision (CV) neemt dit de vorm aan van autoregressieve modellen die het volgende frame/token/pixel voorspellen. Deze benadering verschilt echter op verschillende punten van de menselijke cognitie. Ten eerste beïnvloeden menselijke voorspellingen over de toekomst actief interne cognitieve processen. Ten tweede evalueren mensen van nature de plausibiliteit van voorspellingen over toekomstige toestanden. Op basis van deze capaciteit, en ten derde, door te beoordelen wanneer voorspellingen voldoende zijn, besteden mensen een dynamische hoeveelheid tijd aan het maken van een voorspelling. Dit adaptieve proces is analoog aan Systeem 2-denken in de psychologie. Al deze capaciteiten zijn fundamenteel voor het succes van mensen bij hoogwaardig redeneren en plannen. Om daarom de beperkingen van traditionele autoregressieve modellen die deze mensachtige capaciteiten missen aan te pakken, introduceren we Energy-Based World Models (EBWM). EBWM omvat het trainen van een Energy-Based Model (EBM) om de compatibiliteit van een gegeven context en een voorspelde toekomstige toestand te voorspellen. Hierdoor stelt EBWM modellen in staat om alle drie de aspecten van de menselijke cognitie te bereiken die hierboven zijn beschreven. Bovendien hebben we een variant van de traditionele autoregressieve transformer ontwikkeld, speciaal afgestemd op Energy-Based modellen, genaamd de Energy-Based Transformer (EBT). Onze resultaten tonen aan dat EBWM beter schaalt met data en GPU-uren dan traditionele autoregressieve transformers in CV, en dat EBWM veelbelovende vroege schaalbaarheid biedt in NLP. Als gevolg hiervan biedt deze benadering een spannend pad naar het trainen van toekomstige modellen die in staat zijn tot Systeem 2-denken en intelligente zoekacties in toestandsruimten.

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.