cs.AISep 29, 2026
SaveA neural network that maintains and retrieves memories based on context
Organizations: Department of Psychology, University of Texas at Austin · Center for Theoretical and Computational Neuroscience, Washington University in St. Louis · Department of Neuroscience, University of Texas at Austin · Department of Neuroscience, City University of Hong Kong · Department of Electrical and Systems Engineering, Washington University in St. Louis · Department of Psychology, University of Chicago · Neuroscience Institute, University of Chicago · Institute for Mind and Biology, University of Chicago · Department of Psychology, Washington University in St. Louis
Abstract
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.
Figures & tables
Figure 1: Model description. (A) Context-modulated EM-RNN. The RNN receives video and audio embeddings at each time step and is trained to predict semantics of the next scene. The ground-truth context signals are provided during training, but the model infers context at test via Bayesian inference. Context modulates WM by modulating recurrent connectivity of the RNN in a low-rank manner. The RNN is augmented with an EM buffer. At each step, the input is linearly transformed into a key and a query, and the RNN hidden state serves as a value. Values, keys, and contexts are encoded in the EM. For retrieval, content similarity is computed by comparing the current query and stored keys. Context similarity is computed by comparing the current and past contexts. The EM buffer containing memory content and context is inspired by the hippocampus. The two pattern similarity values (orange) are multiplied element-wise and passed through a softmax ( = 0.1) to produce selective retrieval weights. They are used to compute a weighted sum of stored values, yielding the retrieved memory vector. Such retrieval mechanism of context-congruent memories is inspired by the mPFC. The current scene (RNN hidden state) and the retrieved memory are averaged and linearly transformed to produce the predicted output. The model is trained to minimize the prediction loss. Figure S1 A shows schematics of the original EM-RNN by Song and others (2026) , without context modulation. (B) Movie features. The video frames and audio of each scene are converted to pretrained CLIP and CLAP embedding vectors, respectively, which are reduced to top 50 PCs and used as inputs to the RNN. We annotated which unique characters appeared and where and when the scene took place as binary semantic vectors, reduced them to the top 10 PCs, and used them as RNN outputs. The storyline each scene belongs to represents its context. Each feature dimension of the input and output is -normalized, whereas context signal at each time step sums to 1. (C) Low-rank WM modulation. Four pairs of vectors ( ) are multiplied to represent a unit-rank structured connectivity matrix ( ) for each context. The base recurrence matrix ( ) is then gated (i.e., element-wise multiplied) by the weighted summation of the unit-rank matrices. Figure S1 B shows schematics of the full WM modulation condition.
Figure 2: Context-modulated RNN. (A) Model–brain representational similarity. The scene-by-scene RSMs from each model condition (averaged across 20 random seeds) and each of 200 cortical regions from 33 fMRI participants were correlated, and the mean values were plotted on the lateral and medial surface of the left hemisphere. (Results are comparable across hemispheres so this visualization choice is for simplicity.) (B) Next-scene prediction accuracy at test. Accuracy was measured as the correlation between predicted and observed semantics of the next scene, repeatedly across 100 training iterations. (C) Bayesian posterior probability of the context at test, aligned to moments of ground-truth context shifts. The figure plots results after 0th, 10th, and 100th training iterations. (D) Bayesian context inference compared to the ground truth. The figure shows a segment of example scenes from an example seed. (E) Forward simulation of an example seed. After training, the RNN was forward simulated without input, starting from a zero-vector hidden state and no context signal (black dot). Context signal was then gradually infused toward one of the four contexts in steps of 0.001 until it reached 1 (colored stars). Trajectories from the start to the end are shown for both WM modulation conditions. (F) Tortuosity. The curvature of each trajectory from start to end was quantified, where lower values indicate straighter dynamics. (G) Dimensionality. PCA was applied to the trajectories across all four contexts. The explained variance of the top 20 PCs is shown. (H) Orthogonality. The absolute cosine similarity between all pairs of end points was computed as a measure of separability between context representations, where values near zero indicate greater separability. (B, C, G) Lines indicate the mean and shaded areas indicate standard error of the mean across 20 random seeds. (F, H) Triangles indicate the mean of each distribution.
Figure 3: Performance of context-modulated EM-RNN. (A) Model–brain representational similarity. The scene-by-scene RSMs from each model condition (averaged across 20 random seeds) and each of 200 cortical regions from 33 fMRI participants were correlated, and the mean values were plotted on the lateral and medial surface of the left hemisphere. (B) Next-scene prediction accuracy at test. Accuracy was measured as the correlation between predicted and observed semantics of the next scene, repeatedly across 200 training iterations. (C) Bayesian posterior probability of the context at test, aligned to moments of ground-truth context shifts. The figure plots results after 0th, 10th, and 200th training iterations. A steeper slope following the context shift ( = 0) indicates faster adaptation to a new context. (B-C) Lines indicate the mean and shaded areas indicate standard error of the mean across 20 random seeds.
Figure 4: Memory retrieval of human and EM-RNN. (A) Human memory retrieval. An undirected graph representing the likelihood of one event retrieving another event. This is an average matrix from 33 participants who watched the episode during fMRI. Events 46-48 are omitted in the original study ( Song et al., 2026 ) . (B) Model memory retrieval in the WM+EM modulation condition. Content similarity was computed as a scaled dot product between current query and past key patterns. Context similarity was computed as a scaled dot product between the current and past context probabilities. Content and context similarity values were multiplied and passed through a softmax of , resulting in the final model retrieval matrix on the right. The figure shows the result from an example seed. (C) Model–human retrieval similarity across training iterations. Similarity was computed between matrices A and B right . (D) Similarity between human retrieval matrix and content similarity matrix (solid lines). Similarity was computed between matrices A and B left . A dashed line in the WM+EM modulation condition indicates correlations between human retrieval and the element-wise multiplication of content and context similarity matrices, which was feasible to calculate only in the WM+EM modulation condition because EM buffers in the WM and no modulation conditions do not represent context. (E) Model–human retrieval similarity with context given at test. The analysis is the same as in C but ground-truth context signals are provided at test rather than inferred by the Bayesian observer ( left ). The right panel shows the partial correlation of the model–human retrieval similarity computed after controlling for ground-truth context similarity. (C-E) Lines indicate the mean and shaded areas indicate standard error of the mean across 20 random seeds.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure S2: Comparison of the model–brain RSM similarity between low-rank vs. full WM modulation conditions. Scene-by-scene RSM was computed for each of the 200 cortical regions, respectively in 33 participants’ data. The model’s RSM was computed from the low-rank and full WM modulation conditions respectively, and RSMs across 20 seeds were averaged. The model–brain RSM similarity was computed using Pearson’s correlations, which were Fisher’s -transformed. Thirty-three values were compared between the two model conditions using paired -tests, repeatedly across 200 regions. The figure shows significant cortical regions’ statistics, after FDR-correction of the values. All significant brain regions show positive statistics, meaning the model–brain RSM similarity was higher in the low-rank compared to the full WM modulation condition.
| Input | Low-rank WM | Full WM | Output | None | |
|---|---|---|---|---|---|
| Next-scene prediction accuracy at train ( ) | 0.850 0.006 | 0.801 0.010 | 0.860 0.006 | 0.810 0.010 | 0.772 0.007 |
| Next-scene prediction accuracy at test ( ) | 0.514 0.018 | 0.583 0.019 | 0.555 0.025 | 0.597 0.023 | 0.514 0.017 |
| Bayesian context inference accuracy (%) | 58.73 2.19 | 60.61 5.32 | 53.23 4.15 | 75.40 2.66 | - |
| Model–brain representation similarity ( ) | 0.0039 0.0011 | 0.0087 0.0020 | 0.0080 0.0019 | 0.0053 0.0013 | 0.0066 0.0016 |
Table S1: Comparison of the RNN modulation conditions. Next-scene prediction accuracy is the average Pearson’s correlation between the predicted and observed semantic embeddings of the next scenes, measured either on the training dataset or on the held-out test dataset. Bayesian context inference accuracy is the proportion of scenes in which the inferred context matches the ground-truth context. Model–brain representation similarity is the Pearson’s correlation between fMRI participants’ brain RSMs and the model’s hidden-state RSMs, averaged across 33 participants and 200 cortical regions. All measures were computed after 100 training iterations.
| WM+EM | WM | None | |
|---|---|---|---|
| Next-scene prediction accuracy at train ( ) | 0.799 0.011 | 0.804 0.012 | 0.778 0.008 |
| Next-scene prediction accuracy at test ( ) | 0.565 0.030 | 0.494 0.030 | 0.414 0.043 |
| Bayesian context inference accuracy (%) | 62.10 4.73 | 61.22 4.63 | - |
| Model–brain representation similarity ( ) | 0.0032 0.0009 | 0.0027 0.0009 | 0.0020 0.0007 |
| Model–human retrieval similarity ( ) | 0.283 0.054 | 0.207 0.054 | 0.265 0.039 |
Table S2: Comparison of the EM-RNN modulation conditions. WM modulation was implemented as a low-rank modulation of the RNN’s recurrent connectivity, and EM modulation was implemented by combining context similarity to compute retrieval weights. The measures are defined as in Table S1 . Model–human retrieval similarity is the Pearson’s correlation between the event-by-event (48 48) human and model retrieval matrices.