Modelos de Mundo Baseados em Energia Inspirados na Cognição

Resumo

Um dos métodos predominantes para treinar modelos de mundo é a previsão autoregressiva no espaço de saída do próximo elemento de uma sequência. No Processamento de Linguagem Natural (PLN), isso assume a forma de Modelos de Linguagem de Grande Escala (LLMs) prevendo o próximo token; na Visão Computacional (CV), isso assume a forma de modelos autoregressivos prevendo o próximo quadro/token/pixel. No entanto, essa abordagem difere da cognição humana em vários aspectos. Primeiro, as previsões humanas sobre o futuro influenciam ativamente os processos cognitivos internos. Segundo, os humanos naturalmente avaliam a plausibilidade das previsões em relação a estados futuros. Com base nessa capacidade, e terceiro, ao avaliar quando as previsões são suficientes, os humanos alocam uma quantidade dinâmica de tempo para fazer uma previsão. Esse processo adaptativo é análogo ao pensamento do Sistema 2 na psicologia. Todas essas capacidades são fundamentais para o sucesso dos humanos no raciocínio e planejamento de alto nível. Portanto, para abordar as limitações dos modelos autoregressivos tradicionais que carecem dessas capacidades semelhantes às humanas, introduzimos os Modelos de Mundo Baseados em Energia (EBWM). O EBWM envolve o treinamento de um Modelo Baseado em Energia (EBM) para prever a compatibilidade de um determinado contexto e um estado futuro previsto. Ao fazer isso, o EBWM permite que os modelos alcancem todos os três aspectos da cognição humana descritos. Além disso, desenvolvemos uma variante do transformador autoregressivo tradicional adaptado para modelos baseados em energia, denominado Transformador Baseado em Energia (EBT). Nossos resultados demonstram que o EBWM escala melhor com dados e Horas de GPU do que os transformadores autoregressivos tradicionais em CV, e que o EBWM oferece uma promissora escalabilidade inicial em PLN. Consequentemente, essa abordagem oferece um caminho emocionante para treinar modelos futuros capazes de pensamento do Sistema 2 e de buscar inteligentemente através de espaços 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.

Modelos de Mundo Baseados em Energia Inspirados na Cognição

Cognitively Inspired Energy-Based World Models

Resumo

Support