Модели мира на основе энергии, вдохновленные когнитивными процессами

Cognitively Inspired Energy-Based World Models

June 13, 2024

Авторы: Alexi Gladstone, Ganesh Nanduru, Md Mofijul Islam, Aman Chadha, Jundong Li, Tariq Iqbal

cs.AI

Аннотация

Одним из преобладающих методов обучения мировых моделей является авторегрессионное предсказание в пространстве вывода следующего элемента последовательности. В обработке естественного языка (NLP) это принимает форму больших языковых моделей (LLM), предсказывающих следующий токен; в компьютерном зрении (CV) это принимает форму авторегрессионных моделей, предсказывающих следующий кадр/токен/пиксель. Однако этот подход отличается от когнитивных процессов человека по нескольким аспектам. Во-первых, человеческие предсказания о будущем активно влияют на внутренние когнитивные процессы. Во-вторых, люди естественно оценивают правдоподобие предсказаний относительно будущих состояний. Основываясь на этой способности, и, в-третьих, оценивая, когда предсказания достаточны, люди выделяют динамическое количество времени для предсказания. Этот адаптивный процесс аналогичен мышлению системы 2 в психологии. Все эти способности являются фундаментальными для успеха людей в высокоуровневом рассуждении и планировании. Поэтому, чтобы преодолеть ограничения традиционных авторегрессионных моделей, не обладающих этими человекоподобными способностями, мы представляем Мировые модели на основе энергии (EBWM). EBWM включает в себя обучение модели на основе энергии (EBM) для предсказания совместимости данного контекста и предсказанного будущего состояния. Таким образом, EBWM позволяет моделям достичь всех трех аспектов человеческого когнитивного процесса, описанных выше. Более того, мы разработали вариант традиционного авторегрессионного трансформера, адаптированный для моделей на основе энергии, названный Трансформер на основе энергии (EBT). Наши результаты показывают, что EBWM масштабируется лучше с данными и часами работы GPU, чем традиционные авторегрессионные трансформеры в CV, и что EBWM обещает перспективное раннее масштабирование в NLP. Следовательно, этот подход предлагает захватывающий путь к обучению будущих моделей, способных к мышлению системы 2 и интеллектуальному поиску по пространствам состояний.

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.