# From Memories to Maps: Mechanisms of In-Context Reinforcement Learning in Transformers

arXiv: https://arxiv.org/abs/2506.19686 · HTML: https://mlsanity.com/paper/2506.19686/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 28 of its 61 references that are in the map.
- First posted on arXiv 24 Jun 2025.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [Measuring Compositional Generalization: A Comprehensive Method on Realistic Data](https://mlsanity.com/paper/1912.09713/) (2019)
- [Meta Networks](https://mlsanity.com/paper/1703.00837/) (2017)
- [Using Fast Weights to Attend to the Recent Past](https://mlsanity.com/paper/1610.06258/) (2016)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [Human-Timescale Adaptation in an Open-Ended Task Space](https://mlsanity.com/paper/2301.07608/) (2023)
- [Large Associative Memory Problem in Neurobiology and Machine Learning](https://mlsanity.com/paper/2008.06996/) (2020)
- [Relating transformers to models and neural representations of the hippocampal formation](https://mlsanity.com/paper/2112.04035/) (2021)
- [ICLR: In-Context Learning of Representations](https://mlsanity.com/paper/2501.00070/) (2024)
- [The broader spectrum of in-context learning](https://mlsanity.com/paper/2412.03782/) (2024)
- [(How) Do Language Models Track State?](https://mlsanity.com/paper/2503.02854/) (2025)
- [Attention Approximates Sparse Distributed Memory](https://mlsanity.com/paper/2111.05498/) (2021)
- [Been There, Done That: Meta-Learning with Episodic Recall](https://mlsanity.com/paper/1805.09692/) (2018)
- [Key-value memory in the brain](https://mlsanity.com/paper/2501.02950/) (2025)
- [Rapid Task-Solving in Novel Environments](https://mlsanity.com/paper/2006.03662/) (2020)
- [MEMO: A Deep Network for Flexible Combination of Episodic Memories](https://mlsanity.com/paper/2001.10913/) (2020)
- [Biological learning in key-value memory networks](https://mlsanity.com/paper/2110.13976/) (2021)
- [Graph Structured Prediction Energy Networks](https://mlsanity.com/paper/1910.14670/) (2019)
- [Flexible Prefrontal Control over Hippocampal Episodic Memory for Goal-Directed Generalization](https://mlsanity.com/paper/2503.02303/) (2025)

## Built on by (most cited first, 1 of 1)

- [A neural network that maintains and retrieves memories based on context](https://mlsanity.com/paper/2609.37791/) (2026)

## Abstract (arXiv)

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.