Episodic memory plays a crucial role in various cognitive processes, such as the ability to mentally recall past events. While cognitive science emphasizes the significance of spatial context in the formation and retrieval of episodic memory, the current primary approach to implementing episodic memory in AI systems is through transformers that store temporally ordered experiences, which overlooks the spatial dimension. As a result, it is unclear how the underlying structure could be extended to incorporate the spatial axis beyond temporal order alone and thereby what benefits can be obtained. To address this, this paper explores the use of Spatially-Aware Transformer models that incorporate spatial information. These models enable the creation of place-centric episodic memory that considers both temporal and spatial dimensions. Adopting this approach, we demonstrate that memory utilization efficiency can be improved, leading to enhanced accuracy in various place-centric downstream tasks. Additionally, we propose the Adaptive Memory Allocator, a memory management method based on reinforcement learning that aims to optimize efficiency of memory utilization. Our experiments demonstrate the advantages of our proposed model in various environments and across multiple downstream tasks, including prediction, generation, reasoning, and reinforcement learning. The source code for our models and experiments will be available at https://github.com/junmokane/spatially-aware-transformer.

翻译：情节记忆在各种认知过程中发挥着关键作用，例如在心理上回忆过去事件的能力。尽管认知科学强调空间背景在情节记忆形成与检索中的重要性，但当前在AI系统中实现情节记忆的主要方法是通过存储时序排序经验的Transformer，这忽视了空间维度。因此，尚不清楚底层结构如何在时序顺序之外拓展以融入空间轴，以及由此能获得何种收益。为解决这一问题，本文探索了融合空间信息的空间感知Transformer模型。这些模型能够创建同时考虑时间和空间维度的地点中心情节记忆。采用这一方法，我们证明了可以提高记忆利用效率，从而提升多种地点中心下游任务的准确性。此外，我们提出自适应记忆分配器——一种基于强化学习的记忆管理方法，旨在优化记忆利用效率。我们的实验在多种环境及多个下游任务中（包括预测、生成、推理和强化学习）展示了所提模型的优势。本文模型与实验的源代码将在以下网址公开：https://github.com/junmokane/spatially-aware-transformer。