From memories to maps: Mechanisms of in context reinforcement learning in transformers

2025-06-24

AI Summary

This work addresses the challenge of enabling rapid environmental adaptation beyond the limitations of conventional incremental reinforcement learning. We propose modeling episodic memory as a reusable computational resource and develop a context-aware reinforcement learning framework based on the Transformer architecture, inspired by rodent behavioral planning. Our approach trains agents on biologically grounded navigation tasks requiring contextual generalization. We find that the model spontaneously develops a hybrid planning strategy—blending model-free and model-based characteristics—through cross-contextual representation alignment and implicit structural learning; crucially, its memory stores not only experiences but also intermediate computational states to support flexible decision-making. Neural representational analyses reveal that internal activation patterns closely mirror functional properties of the hippocampal–entorhinal circuit, providing an interpretable AI model for investigating the computational principles underlying biological episodic memory.

Technology Category

• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Machine Learning: Reinforcement Learning
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Search and Retrieval-Augmented AI: Agentic search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Psychology-informed user models and recommender systems

Abstract

Humans and animals show remarkable learning efficiency, adapting to new environments with minimal experience. This capability is not well captured by standard reinforcement learning algorithms that rely on incremental value updates. Rapid adaptation likely depends on episodic memory -- the ability to retrieve specific past experiences to guide decisions in novel contexts. Transformers provide a useful setting for studying these questions because of their ability to learn rapidly in-context and because their key-value architecture resembles episodic memory systems in the brain. We train a transformer to in-context reinforcement learn in a distribution of planning tasks inspired by rodent behavior. We then characterize the learning algorithms that emerge in the model. We first find that representation learning is supported by in-context structure learning and cross-context alignment, where representations are aligned across environments with different sensory stimuli. We next demonstrate that the reinforcement learning strategies developed by the model are not interpretable as standard model-free or model-based planning. Instead, we show that in-context reinforcement learning is supported by caching intermediate computations within the model's memory tokens, which are then accessed at decision time. Overall, we find that memory may serve as a computational resource, storing both raw experience and cached computations to support flexible behavior. Furthermore, the representations developed in the model resemble computations associated with the hippocampal-entorhinal system in the brain, suggesting that our findings may be relevant for natural cognition. Taken together, our work offers a mechanistic hypothesis for the rapid adaptation that underlies in-context learning in artificial and natural settings.

Problem

• Mechanisms of rapid adaptation in transformers
• In-context reinforcement learning in planning tasks
• Memory as computational resource for flexible behavior

Innovation

• Transformers enable in-context reinforcement learning
• Memory tokens cache intermediate computations
• Representations align across diverse sensory environments