# A neural network that maintains and retrieves memories based on context

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 63 references that are in the map.
- First posted on arXiv 29 Sep 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Large-scale Contrastive Language-Audio Pretraining with Feature Fusion and Keyword-to-Caption Augmentation](https://mlsanity.com/paper/2211.06687/) (2022)
- [Neural Episodic Control](https://mlsanity.com/paper/1703.01988/) (2017)
- [Large Associative Memory Problem in Neurobiology and Machine Learning](https://mlsanity.com/paper/2008.06996/) (2020)
- [Linking connectivity, dynamics and computations in low-rank recurrent neural networks](https://mlsanity.com/paper/1711.09672/) (2017)
- [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)
- [Flexible Prefrontal Control over Hippocampal Episodic Memory for Goal-Directed Generalization](https://mlsanity.com/paper/2503.02303/) (2025)
- [From Memories to Maps: Mechanisms of In-Context Reinforcement Learning in Transformers](https://mlsanity.com/paper/2506.19686/) (2025)

## Abstract (arXiv)

Every day, people continuously infer situational context and adjust the way they understand and remember the world. Context, signaled by the prefrontal cortex, is known to modulate working memory and episodic memory, but the algorithmic understanding of this modulation remains limited. Here, we train a recurrent neural network (RNN), augmented with an episodic memory buffer, to infer context using Bayesian inference as it continuously makes predictions of upcoming scenes while watching naturalistic movies. When the inferred context modulates the RNN's recurrent connectivity (the basis of working memory) in a low-rank manner, the model's activity patterns best match neural responses in human participants who watched the same movies during fMRI. Context also modulates episodic memory retrieval, such that the model retrieves memories based on not only content similarity but also context similarity. This is implemented as a key-value system with self-attention, designed to additionally encode context and retrieve context-congruent memories. The resulting model not only better resembles human brain representations but also learns to retrieve memories like humans much faster than a model without context modulation. Together, our findings suggest a computational mechanism by which context modulates information maintenance and long-term memory retrieval in naturalistic environments.