🤖 AI Summary
Current AI systems lack biologically inspired cognitive map capabilities, hindering strong out-of-distribution (OOD) generalization and flexible path integration. To address this, we propose MapFormers: a Transformer-based self-supervised framework for cognitive map modeling. Its core innovation is input-dependent dynamic positional encoding, which naturally disentangles spatial structure from semantic content while unifying absolute (episodic) and relative (working) positional representations. This design enables parallel path integration and faithful simulation of memory systems. Evaluated on 2D navigation tasks, MapFormers achieves near-perfect OOD generalization—substantially outperforming state-of-the-art methods—and empirically validates the critical role of structural priors in cognitive map learning.
📝 Abstract
A cognitive map is an internal model which encodes the abstract relationships among entities in the world, giving humans and animals the flexibility to adapt to new situations, with a strong out-of-distribution (OOD) generalization that current AI systems still do not possess. To bridge this gap, we introduce MapFormers, new architectures based on Transformer models, which can learn cognitive maps from observational data and perform path integration in parallel, in a self-supervised manner. Cognitive maps are learned in the model by disentangling structural relationships in the inputs from their specific content, a property that can be achieved naturally by updating the positional encoding in Transformers with input-dependent matrices. We developed two variants of MapFormers that unify absolute and relative positional encoding to model episodic (EM) and working memory (WM), respectively. We tested MapFormers on several tasks, including a classic 2D navigation task, showing that our models can learn a cognitive map of the underlying space and generalize OOD (e.g., to longer sequences) with near-perfect performance, unlike current architectures. Together, these results demonstrate the superiority of models designed to learn a cognitive map, and the importance of introducing a structural bias for structure-content disentanglement, which can be achieved in Transformers with input-dependent positional encoding. MapFormers have broad applications in both neuroscience and AI, by explaining the neural mechanisms giving rise to cognitive maps, while allowing these relation models to be learned at scale.