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

TL;DR

This paper investigates how transformers can emulate episodic memory to enable rapid, in-context reinforcement learning, revealing that memory caching supports flexible decision-making akin to hippocampal functions.

Contribution

It introduces a transformer-based model that uses memory caching for in-context reinforcement learning, offering a mechanistic understanding of rapid adaptation and representation alignment.

Findings

Memory supports in-context learning through caching computations.

Representations align across different environments, resembling hippocampal functions.

In-context RL strategies differ from standard model-free or model-based approaches.

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.