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

• URL: http://arxiv.org/abs/2506.19686v2
• Date: Thu, 26 Jun 2025 17:18:54 GMT
• Title: From Memories to Maps: Mechanisms of In-Context Reinforcement Learning in Transformers
• Authors: Ching Fang, Kanaka Rajan,
• Abstract summary: We train a transformer to in-context reinforcement learn in a distribution of planning tasks inspired by rodent behavior.<n>We characterize the learning algorithms that emerge in the model.<n>We find that memory may serve as a computational resource, storing both raw experience and cached computations to support flexible behavior.
• Score: 2.4554686192257424
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• A Neural Network Model of Complementary Learning Systems: Pattern Separation and Completion for Continual Learning [2.9123921488295768]
  Learning new information without forgetting prior knowledge is central to human intelligence.<n>In contrast, neural network models suffer from catastrophic forgetting when acquiring new information.<n>We develop a neurally plausible continual learning model that achieves close to state-of-the-art accuracy (90%)<n>Our work provides a functional template for modeling memory consolidation, generalization, and continual learning in both biological and artificial systems.
  arXiv  Detail & Related papers  (2025-07-15T15:05:26Z)
• Sequence-to-Sequence Models with Attention Mechanistically Map to the Architecture of Human Memory Search [13.961239165301315]
  We show that foundational architectures in neural machine translation exhibit mechanisms that directly correspond to those specified in the Context Maintenance and Retrieval model of human memory.<n>We implement a neural machine translation model as a cognitive model of human memory search that is both interpretable and capable of capturing complex dynamics of learning.
  arXiv  Detail & Related papers  (2025-06-20T18:43:15Z)
• Rethinking Associative Memory Mechanism in Induction Head [37.93644115914534]
  This paper investigates how a two-layer transformer thoroughly captures in-context information and balances it with pretrained bigram knowledge in next token prediction.<n>We theoretically analyze the representation of weight matrices in attention layers and the resulting logits when a transformer is given prompts generated by a bigram model.
  arXiv  Detail & Related papers  (2024-12-16T05:33:05Z)
• In-Memory Learning: A Declarative Learning Framework for Large Language Models [56.62616975119192]
  We propose a novel learning framework that allows agents to align with their environment without relying on human-labeled data. This entire process transpires within the memory components and is implemented through natural language. We demonstrate the effectiveness of our framework and provide insights into this problem.
  arXiv  Detail & Related papers  (2024-03-05T08:25:11Z)
• Revisiting Dynamic Evaluation: Online Adaptation for Large Language Models [88.47454470043552]
  We consider the problem of online fine tuning the parameters of a language model at test time, also known as dynamic evaluation. Online adaptation turns parameters into temporally changing states and provides a form of context-length extension with memory in weights.
  arXiv  Detail & Related papers  (2024-03-03T14:03:48Z)
• Spatially-Aware Transformer for Embodied Agents [20.498778205143477]
  This paper explores the use of Spatially-Aware Transformer models that incorporate spatial information. We demonstrate that memory utilization efficiency can be improved, leading to enhanced accuracy in various place-centric downstream tasks. We also propose the Adaptive Memory Allocator, a memory management method based on reinforcement learning.
  arXiv  Detail & Related papers  (2024-02-23T07:46:30Z)
• Evolving Reservoirs for Meta Reinforcement Learning [1.6874375111244329]
  We propose a computational model for studying a mechanism that can enable such a process. At the evolutionary scale, we evolve reservoirs, a family of recurrent neural networks. We employ these evolved reservoirs to facilitate the learning of a behavioral policy through Reinforcement Learning (RL) Our results show that the evolution of reservoirs can improve the learning of diverse challenging tasks.
  arXiv  Detail & Related papers  (2023-12-09T16:11:48Z)
• Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability [39.42414275888214]
  We evaluate 13 models capable of causal language modeling across a suite of synthetic in-context learning tasks. All the considered architectures can perform in-context learning under a wider range of conditions than previously documented. Several attention alternatives are sometimes competitive with or better in-context learners than transformers.
  arXiv  Detail & Related papers  (2023-10-12T05:43:06Z)
• What Can Transformers Learn In-Context? A Case Study of Simple Function Classes [67.06980111346245]
  In-context learning refers to the ability of a model to condition on a prompt sequence consisting of in-context examples. We show that standard Transformers can be trained from scratch to perform in-context learning of linear functions. We also show that we can train Transformers to in-context learn more complex function classes with performance that matches or exceeds task-specific learning algorithms.
  arXiv  Detail & Related papers  (2022-08-01T18:01:40Z)
• Decoupling Knowledge from Memorization: Retrieval-augmented Prompt Learning [113.58691755215663]
  We develop RetroPrompt to help a model strike a balance between generalization and memorization. In contrast with vanilla prompt learning, RetroPrompt constructs an open-book knowledge-store from training instances. Extensive experiments demonstrate that RetroPrompt can obtain better performance in both few-shot and zero-shot settings.
  arXiv  Detail & Related papers  (2022-05-29T16:07:30Z)
• Representation Learning for Context-Dependent Decision-Making [22.16801879707937]
  We study representation learning in the sequential decision-making scenario with contextual changes. We propose an online algorithm that is able to learn and transfer context-dependent representations.
  arXiv  Detail & Related papers  (2022-05-12T01:06:57Z)
• Dynamic Inference with Neural Interpreters [72.90231306252007]
  We present Neural Interpreters, an architecture that factorizes inference in a self-attention network as a system of modules. inputs to the model are routed through a sequence of functions in a way that is end-to-end learned. We show that Neural Interpreters perform on par with the vision transformer using fewer parameters, while being transferrable to a new task in a sample efficient manner.
  arXiv  Detail & Related papers  (2021-10-12T23:22:45Z)
• Backprop-Free Reinforcement Learning with Active Neural Generative Coding [84.11376568625353]
  We propose a computational framework for learning action-driven generative models without backpropagation of errors (backprop) in dynamic environments. We develop an intelligent agent that operates even with sparse rewards, drawing inspiration from the cognitive theory of planning as inference. The robust performance of our agent offers promising evidence that a backprop-free approach for neural inference and learning can drive goal-directed behavior.
  arXiv  Detail & Related papers  (2021-07-10T19:02:27Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.