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.

通过训练基于能量的世界模型，我们引入了一种能够模拟人类认知的方法，其通过预测上下文和未来状态的兼容性来解决传统的自回归模型的局限性，并展示了在计算机视觉和自然语言处理中的潜在应用。