认知启发的基于能量的世界模型

训练世界模型的主要方法之一是在输出空间中对序列的下一个元素进行自回归预测。在自然语言处理（NLP）中，这表现为大型语言模型（LLM）预测下一个词元；在计算机视觉（CV）中，这表现为自回归模型预测下一帧/词元/像素。然而，这种方法在几个方面与人类认知不同。首先，人类对未来的预测会主动影响内部认知过程。其次，人类自然地评估关于未来状态的预测的合理性。基于这种能力，第三，通过评估预测何时足够，人类会分配动态的时间来进行预测。这种自适应过程类似于心理学中的系统2思维。所有这些能力都是人类在高级推理和规划中取得成功的基础。因此，为了解决传统自回归模型缺乏这些类人能力的局限性，我们引入了基于能量的世界模型（EBWM）。EBWM涉及训练一个基于能量的模型（EBM）来预测给定上下文和预测的未来状态之间的兼容性。通过这样做，EBWM使模型能够实现所描述的人类认知的所有三个方面。此外，我们开发了一种针对基于能量的模型量身定制的传统自回归Transformer的变体，称为基于能量的Transformer（EBT）。我们的结果表明，在CV中，EBWM在数据和GPU时间上的扩展性优于传统自回归Transformer，并且在NLP中EBWM展现出有前景的早期扩展性。因此，这种方法为训练未来能够进行系统2思维并在状态空间中进行智能搜索的模型提供了一条令人兴奋的路径。