인지 과학에서 영감을 받은 에너지 기반 세계 모델

Cognitively Inspired Energy-Based World Models

June 13, 2024

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

cs.AI

초록

세계 모델(world model)을 학습시키는 주요 방법 중 하나는 시퀀스의 다음 요소를 출력 공간에서 자기회귀적으로 예측하는 것입니다. 자연어 처리(NLP)에서는 대형 언어 모델(LLM)이 다음 토큰을 예측하는 형태로 나타나며, 컴퓨터 비전(CV)에서는 자기회귀 모델이 다음 프레임/토큰/픽셀을 예측하는 형태로 나타납니다. 그러나 이 접근 방식은 인간의 인지와 몇 가지 측면에서 차이가 있습니다. 첫째, 인간의 미래 예측은 내부 인지 과정에 적극적으로 영향을 미칩니다. 둘째, 인간은 미래 상태에 대한 예측의 타당성을 자연스럽게 평가합니다. 이 능력을 바탕으로, 셋째, 예측이 충분한지 평가함으로써 인간은 예측을 위해 동적으로 시간을 할당합니다. 이 적응적 과정은 심리학에서의 시스템 2 사고(System 2 thinking)와 유사합니다. 이러한 모든 능력은 인간이 고수준의 추론과 계획을 성공적으로 수행하는 데 근본적입니다. 따라서 이러한 인간과 유사한 능력이 부족한 전통적인 자기회귀 모델의 한계를 해결하기 위해, 우리는 에너지 기반 세계 모델(Energy-Based World Models, EBWM)을 제안합니다. EBWM은 주어진 맥락과 예측된 미래 상태의 호환성을 예측하도록 에너지 기반 모델(Energy-Based Model, EBM)을 학습시킵니다. 이를 통해 EBWM은 앞서 설명한 인간 인지의 세 가지 측면을 모두 달성할 수 있게 합니다. 더 나아가, 우리는 에너지 기반 모델에 맞게 조정된 전통적인 자기회귀 트랜스포머의 변형인 에너지 기반 트랜스포머(Energy-Based Transformer, EBT)를 개발했습니다. 우리의 실험 결과는 EBWM이 CV에서 전통적인 자기회귀 트랜스포머보다 데이터와 GPU 시간에 대해 더 나은 확장성을 보이며, 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.