🤖 machine learning

MapFormer: Self-Supervised Learning of Cognitive Maps with Input-Dependent Positional Embeddings

The paper introduces MapFormer, a novel Transformer-based architecture that utilizes input-dependent positional embeddings derived from Lie-algebra generators to learn unsupervised cognitive maps, thereby achieving superior out-of-distribution generalization and scalability across diverse cognitive tasks compared to existing AI models.

Original authors: Victor Rambaud, Salvador Mascarenhas, Yair Lakretz

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Victor Rambaud, Salvador Mascarenhas, Yair Lakretz

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are walking through a new city. You don't just memorize every single building you see; instead, you build a mental map. You understand that "turning left" always changes your position relative to the street, regardless of whether you are in Paris or Tokyo. This ability to separate where you are (the structure) from what you see (the content) is what scientists call a "cognitive map."

Current AI systems, like the large language models you might chat with, are incredibly good at memorizing patterns. However, they are often "brittle." If you ask them to solve a problem slightly different from what they trained on, they often fail. They struggle to understand the underlying rules of a situation and instead rely on guessing based on surface-level similarities.

This paper introduces a new type of AI architecture called MapFormers. Think of MapFormers as giving AI a "mental GPS" that learns the rules of the road rather than just memorizing the scenery.

The Core Idea: Separating the "Where" from the "What"

The authors argue that to be truly smart, an AI needs to untangle two things:

Content: The specific things it sees (e.g., a red apple, a blue car).
Structure: The rules of how things move or relate (e.g., "moving right" always shifts you one step to the right).

Standard AI models mix these together. MapFormers are designed to keep them separate. They do this using a mathematical trick involving Lie groups (a fancy way of describing smooth rotations and movements). Instead of hard-coding rules, the AI learns to generate "movement matrices" based on the input.

The Two Types of MapFormers

The paper presents two versions of this system, which the authors compare to human memory types:

MapEM (Episodic Memory): This is like a diary. It keeps a separate, dedicated list of "where I am" and "what I saw." It's very good at recalling specific past events (like remembering exactly what you had for breakfast last Tuesday), but it's a bit slower to process.
MapWM (Working Memory): This is like your brain's active scratchpad. It mixes the "where" and "what" together on the fly. It's faster and more efficient, similar to how RoPE (a popular current AI technique) works, but with a crucial upgrade: it can actually learn to move around a map dynamically.

How They Learned (The Training Grounds)

To prove these models work, the researchers tested them on three specific "video game" style challenges where the rules were hidden:

The "Ignore the Noise" Game: The AI had to copy a list of items but ignore the blank spaces in between. Standard AI gets confused by the blanks; MapFormers learned to "gate" (block out) the noise and focus only on the items to copy.
The "Blindfolded Navigator" Game: The AI was given a sequence of "move" and "see" commands but wasn't told which was which. It had to figure out that "Up" moves you up and "See a Tree" just updates your view. Once it learned the map, it could predict exactly what it would see if it returned to a spot it had visited before, even if the sequence was much longer than it had ever seen before.
The "Nested Brackets" Game: This tests the ability to handle deep layers of logic (like ((()))). Standard AI struggles to count how many open brackets are waiting to be closed when the sequence gets long. MapFormers treated opening a bracket as "moving forward" and closing it as "moving backward," allowing them to track the stack perfectly, even in sequences much deeper than they were trained on.

The Results: Why This Matters

The paper claims that MapFormers achieved near-perfect generalization.

The "OOD" Test: In AI terms, "Out-of-Distribution" (OOD) means testing the model on scenarios it has never seen before. While standard models failed miserably when the sequences got longer or the environments changed, MapFormers succeeded. They didn't just memorize; they learned the geometry of the problem.
Real-World Test: The researchers also tried this on a chunk of real internet text (OpenWebText). While the improvement was smaller than on the logic games, the MapFormers still performed slightly better than standard models, suggesting these principles might scale up to real-world language tasks.

The "Secret Sauce": Math as a Compass

The paper explains that MapFormers use a specific mathematical approach where actions (like "move right") are treated as rotations.

Imagine you are spinning a dial. If you turn it 90 degrees, then another 90 degrees, you end up at 180.
MapFormers learn that "Right" is a specific rotation and "Left" is the opposite rotation.
Because they use this mathematical framework, they can calculate the result of a long journey (path integration) by simply adding up the angles, rather than doing complex, slow calculations for every single step. This allows them to process information in parallel (very fast) while still understanding the sequence.

Summary

In short, this paper proposes that to make AI more robust and adaptable, we should stop treating it like a giant pattern-matching machine and start giving it a cognitive map. By teaching the AI to separate the rules of movement from the things it sees, MapFormers can navigate new, unseen situations with a flexibility that current AI systems lack. They don't just guess; they understand the map.

Technical Summary: MapFormers

Problem Statement

Current Transformer-based AI systems, despite impressive performance on in-distribution tasks, exhibit significant brittleness when faced with out-of-distribution (OOD) scenarios. They often fail to generalize to novel situations that humans and animals handle with ease. The authors attribute this gap to a lack of structural biases that allow biological brains to construct cognitive maps: internal models that disentangle the abstract, invariant structure of an environment from its specific, variable content.

Existing attempts to model cognitive maps, such as the Tolman-Eichenbaum Machine (TEM) using action-dependent matrices in RNNs, suffer from two critical limitations:

Lack of Scalability: They require sequential matrix multiplication to update positions, preventing parallel processing.
Pre-specified Actions: They often require the set of actions (e.g., cardinal directions) to be hardcoded rather than learned from data.

Furthermore, standard positional encodings (like RoPE) are static or rely on heuristics that do not support the dynamic composition of actions required for path integration in unknown environments.

Methodology

The paper introduces MapFormers, a class of Transformer architectures designed to learn cognitive maps and perform path integration via next-token prediction without supervision. The core innovation lies in using input-dependent matrices to update positional encodings, derived from Lie-group theory.

Theoretical Framework

Lie-Group Formalism: The authors model actions as elements of a Lie group (specifically compact groups like $SO(2)$ for rotations). Instead of learning a separate matrix for every possible action, the model learns Lie-algebra generators.
Path Integration via Exponential Maps: Actions are represented as infinitesimal generators in the Lie algebra. Path integration (composing a sequence of actions) is performed by summing these generators in the algebra (linear accumulation) and then applying a single exponential map to obtain the finite transformation in the Lie group.
- For commutative maps (e.g., 2D grid navigation), this allows path integration to be computed in parallel via cumulative summation, followed by a single matrix exponential.
- This approach reduces the inductive bias from the number of actions to the structural complexity of the algebra.

Architectural Variants

The paper proposes two variants of MapFormers, unifying absolute and relative positional encoding with distinct memory models:

MapEM (Episodic Memory):
- Uses absolute positional embeddings to represent a cognitive map.
- Disentangles content (observations) and structure (position) into separate neural populations.
- Computes attention on the conjunction of content and position ( $Q_G = \text{vec}(Q^T \cdot P)$ ), effectively acting as an attention mask where position encodings update independently of content.
- Suitable for tasks requiring robust recall of specific states.
MapWM (Working Memory):
- Uses relative positional encoding (analogous to RoPE but dynamic).
- Represents structure implicitly by rotating neural activity (keys and queries) into distinct subspaces via input-dependent rotation matrices ( $R_{\theta PI}$ ).
- The rotation angle $\theta$ is derived from a low-rank projection of the input token, allowing the model to learn which tokens act as actions (updating position) and which are observations (leaving position unchanged).

Input-Dependent Mechanism

Both architectures employ a two-stage projection to determine the "integration duration" ( $\Delta_t$ ) for each token:

Inner Projection: Maps the input token to a low-dimensional action representation.
Outer Projection: Projects this representation to a per-block integration duration.
The final rotation angle is $\theta_t = \omega \odot \Delta_t$ , where $\omega$ is a learnable angular velocity vector. Path integration is computed as the cumulative sum of these angles ( $\theta_{0 \to t} = \text{cumsum}(\theta_t)$ ).

Key Contributions

MapFormers Architecture: A novel Transformer variant that learns cognitive maps by disentangling structure from content using input-dependent positional encodings based on Lie-algebra generators.
Parallel Path Integration: By leveraging commutative Lie groups (e.g., $SO(2)^k$ ), the models achieve parallel path integration, overcoming the sequential bottleneck of previous RNN-based cognitive map models.
Unification of Memory Models: The work provides a theoretical link between absolute/relative positional encoding and episodic/working memory, showing how both can be derived from the same Lie-theoretic principles.
Scalability to Non-Commutative Maps: The framework is extended to non-commutative problems (e.g., family trees) using multiple generators, though this requires sequential path integration.

Experimental Results

The authors evaluated MapFormers on three formal tasks targeting distinct cognitive capacities:

Selective Copy (Gating):
- Task: Copy a sequence while ignoring distractor tokens.
- Result: MapFormers (both EM and WM) achieved perfect OOD generalization, whereas standard baselines (RoPE, TAPE) failed. This demonstrates the ability to learn dynamic gating mechanisms.
Grid Navigation (1D and 2D):
- Task: Predict observations in a sequence of alternating actions and random observations, requiring the model to infer the latent topology and path-integrate position.
- Result: MapFormers achieved near-perfect accuracy on OOD sequences (varying length and density), while baselines (including CoPE and PathAtt) failed to generalize.
- Neural Analysis: The models learned to project observation tokens to zero rotation (identity) and action tokens to specific rotation directions, effectively disentangling content from structure.
Dyck-2 (Nested Hierarchies):
- Task: Predict valid continuations of bracket sequences, requiring the tracking of a stack.
- Result: MapFormers generalized perfectly to longer and deeper sequences than seen in training using only a single layer. Baselines required multiple layers and still failed on OOD data.
- Mechanism: The models embedded the stack directly into the positional encoding by using orthogonal dimensions for different bracket types and opposite directions for open/close brackets.
Natural Language Scaling:
- Task: Pretraining on OpenWebText.
- Result: A 12-layer MapWM model showed consistent perplexity improvements over a RoPE baseline ($18.79$ vs $19.14$) and better length extrapolation. However, on the BLIMP syntactic benchmark, MapWM did not show significant gains over RoPE, suggesting that commutative path integration alone may not be sufficient for all syntactic modeling.

Significance and Claims

The paper claims that MapFormers provide a critical structural bias that enables systematic OOD generalization by capturing latent structures via path integration rather than relying on heuristics.

Robustness: The models demonstrate that disentangling abstract relations from content allows for robust generalization to unseen sequence lengths, depths, and environmental densities.
Scalability: The use of parallel computation on commutative maps allows these principles to scale to large-scale, real-world domains, as evidenced by the natural language results.
Biological Plausibility: The approach bridges the gap between biological cognitive maps (which factorize "what" from "where") and modern deep learning, offering a pathway to AI systems that can reason through novel problems by understanding the underlying geometry of their environment.

The authors remain modest regarding the limitations, noting that non-commutative path integration is not parallelizable and that the current commutative constraint may limit the modeling of certain complex syntactic structures in language. They suggest that future work could explore hybrid architectures or non-commutative extensions for more complex relational reasoning.