Source-linked AI summary
A neural network that maintains and retrieves memories based on context
Hayoung Song, JeongJun Park, Qihong Lu, Giacomo Vedovati, Monica D. Rosenberg, Zachariah M. Reagh, ShiNung Ching
TL;DR
Algorithms for how context jointly modulates working and episodic memory remain unclear. The paper trains a recurrent network with an episodic-memory buffer to infer context from naturalistic movie scenes and modulate processing and retrieval. Low-rank recurrent modulation best matched human brain representations, while explicit context in memory retrieval accelerated human-like retrieval when retrieval was selective.
Problem
How context jointly modulates working and episodic memory remains algorithmically unclear, with prior models treating them in separate lines of research.
Source: §1 · p. 2Method
The authors augment a predictive recurrent network with Bayesian context inference that modulates recurrent processing and context-aware key-value episodic-memory retrieval.
Source: §1 · p. 2, §2.2 · p. 4Results
Low-rank working-memory modulation best matched human brain representations, and context-modulated retrieval reached r = 0.2 by the 22nd iteration versus the 94th for the no modulation model.
Source: §3.1 · p. 5, §4.2 · p. 8, §5 · p. 10Takeaways & Limitations
The context-modulated EM-RNN is a candidate mechanism for how context modulates information processing and memory retrieval in naturalistic environments.
Source: §1 · p. 2Model–brain similarity may reflect context modulation of sensory processing regions rather than PFC context representations, because PFC function is implemented computationally rather than as an explicit network module.
Source: §3.1 · p. 5
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Context is known to shape working and episodic memory, but the algorithms that jointly support context-dependent maintenance and retrieval remain unclear. This study tests an inferred-context model against human behavioral and brain data during naturalistic movie viewing.
1 INTRODUCTION
- No model had inferred context to jointly modulate working memory and episodic memory within one framework.Prior work modeled the two memory systems separately, and some working-memory models received context as an explicit input.Source: §1 · p. 2
- The authors compare candidate models with behavioral and fMRI data from people watching the same television episode.Naturalistic movies provide multimodal stimuli, long-range dependencies, and structured storylines.Source: §1 · p. 2
- The proposed system uses Bayesian inference to infer context, then modulates both information maintenance and episodic-memory encoding and retrieval.It builds on a key-value EM-RNN and uses the posterior as the context signal.Source: §1 · p. 2
- Low-rank context modulation of recurrent connectivity best matched human brain representations among the tested working-memory modulation stages.The paper also reports faster human-like memory retrieval when context is explicitly represented and retrieval is selective.Source: §1 · p. 2
2 MODEL
The model watches scenes, predicts upcoming semantics, and uses inferred storyline context to modulate recurrent processing and episodic-memory retrieval. Context is inferred from next-scene prediction, with the posterior serving as the context signal.
2.1 A RECURRENT NEURAL NETWORK WITH AN EPISODIC MEMORY BUFFER THAT WATCHES
- The model receives video and audio scene embeddings and predicts the semantics of the next scene using an RNN augmented with an episodic-memory buffer.The value is the RNN hidden state; keys and values are stored, and retrieval integrates relevant memories with the current scene.Source: §2.1 · pp. 2, 4, Figure 1 · p. 3
2.2 CONTEXT INFERENCE AND MODULATION
- Context modulates recurrent connectivity in a low-rank manner and also modulates memory retrieval using content and context similarity.The model asks how context modulates both the RNN and retrieval from the episodic-memory buffer.Source: Figure 1 · p. 3, §2.2 · p. 4
2.1 A RECURRENT NEURAL NETWORK WITH AN EPISODIC MEMORY BUFFER THAT WATCHES
- Scenes are represented with pretrained visual and auditory embeddings, while output semantics annotate characters and where and when scenes occur.The RNN trains on episodes 2–18 and is tested on episode 1.Source: §2.1 · p. 3
- The episodic-memory system uses separate keys, queries, and values: queries survey stored keys by content similarity, and associated values are retrieved as a weighted average.The retrieved memory is averaged with the current hidden state to predict next-scene semantics.Source: §2.1 · p. 4
2.2 CONTEXT INFERENCE AND MODULATION
- At test, the model infers each scene's context by comparing next-scene predictions under four storylines, using the previous posterior as a persistence-promoting prior.The resulting posterior serves as the context signal; the inference succeeds only if learned parameters make predictions more accurate under the correct context.Source: §2.2 · p. 4
3 CONTEXT MODULATION OF THE RECURRENT NEURAL NETWORK
Low-rank WM modulation yielded stronger model–brain representational similarity than full WM modulation, and its context trajectories were straighter, lower-dimensional, and more separated.
3.1 CONTEXT MODULATES WORKING MEMORY THROUGH LOW-RANK GATING OF RECURRENT
- The comparison tested input, output, fully context-specific recurrence, and low-rank modulation of shared recurrence against an unmodulated baseline.Source: §3.1 · pp. 4–5
- Only WM modulation conditions exceeded the no-modulation baseline, and low-rank WM modulation aligned better with brain representations than full WM modulation.The comparison used scene-by-scene representational similarity matrices across 33 participants and 200 cortical regions.Source: §3.1 · p. 5, Figure 2 · p. 6
- Model–brain similarity was high in visual and auditory regions, so the match reflects sensory-region representations rather than an explicit PFC activity module.Source: §3.1 · p. 5
- At test, next-scene prediction reached r = 0.583 ± 0.019, while Bayesian context inference reached 60.61 ± 5.32%.Source: §3.1 · p. 5, Figure 2 · p. 6
3.2 MODULATING LOW-RANK CONNECTIVITY AFFORDS EFFICIENT DYNAMICS AND CLEAR SEPARATION
- Low-rank WM trajectories were straighter and lower-dimensional than full-WM trajectories, with more orthogonal context endpoints and recurrent connectivity.Simulations began from a zero hidden state with no context signal, which was infused in 0.001 steps; Figure 2E–H shows trajectory and separation measures.Source: §3.2 · pp. 5–7, Figure 2 · p. 6
4 CONTEXT MODULATION OF THE EPISODIC MEMORY RETRIEVAL
The EM buffer combines content and context similarity to retrieve context-congruent memories. This context-modulated system reaches human-like retrieval faster than a model without modulation, provided retrieval is selective.
4.1 AN EPISODIC MEMORY SYSTEM THAT RETRIEVES BASED ON BOTH CONTENT AND CONTEXT
- The EM buffer stores context alongside keys and values, then multiplies content and context similarities to favor memories matching both.The resulting weights are passed through a softmax before high-weight values are retrieved.Source: §4.1 · p. 7, §4.2 · p. 7
- Next-scene prediction correlations rose from r = 0.414 ± 0.043 without modulation to r = 0.565 ± 0.030 with WM+EM modulation.WM+EM also aligned hidden-state representations with brain regions better than WM or no modulation, and shifted context faster than WM despite comparable overall inference accuracy.Source: §4.1 · p. 7, Figure 3 · p. 8
4.2 CONTEXT-MODULATED SYSTEM LEARNS FASTER TO RETRIEVE LIKE HUMANS ONLY WHEN RETRIEVAL IS SELECTIVE
- 22 training iterations reached r = 0.2 in the WM+EM model, versus 94 iterations for the no-modulation model.At the 200th iteration, their model–human retrieval similarities were comparable: r = 0.283 ± 0.054 and r = 0.265 ± 0.039, respectively.Source: §4.2 · p. 8, Figure 4 · p. 9
- Content similarity alone aligned less with human retrieval in WM+EM than without modulation, but combining content and context similarity raised alignment to the WM level.The combined correlation was r = 0.274 ± 0.071, still below the no-modulation condition.Source: §4.2 · p. 8, Figure 4 · p. 9
- Context modulation improved human-like retrieval only when retrieval was selective; increasing softmax temperature reduced the WM+EM model’s resemblance to human retrieval.Human retrieval was selective, supporting the relevance of a small temperature.Source: §4.2 · pp. 8–9
- At test, retrieval performance depended on Bayesian context inference; providing ground-truth context increased model–human similarity before training, while similarity still grew after controlling for context similarity.Context inference accuracy rose sharply during the first ∼30 iterations and then stabilized, paralleling the early increase in retrieval similarity.Source: §4.2 · pp. 9–10, Figure 4 · p. 9
- Modulating input-to-key and input-to-query transformations instead of retrieval produced lower model–human retrieval similarity than WM+EM modulation.The alternative achieved r = 0.139 ± 0.069, compared with r = 0.283 ± 0.054 for WM+EM modulation.Source: §4.2 · p. 10
5 CONCLUSION
The conclusion identifies context modulation of working and episodic memory as a candidate computational account, based on comparisons with naturalistic behavioral and fMRI data.
5 CONCLUSION
- The context-modulated RNN best matched human brain representations when context modulated working memory, and its context-modulated episodic memory retrieved like humans faster when retrieval was selective.Structured low-rank modulation of shared recurrent connectivity yielded more efficient dynamics and clearer separation between contexts than independent context-specific connectivity.Source: §5 · p. 10
- The authors report using generative AI tools for readability, LaTeX conversion, equation checks, and figure-generation code, with authors retaining responsibility for the final work.The research ideas, model and analysis code, experiments, and manuscript draft are identified as the authors’ original work.Source: AI USE STATEMENT · p. 10
- The paper’s code and the cited study’s fMRI and behavioral data are available through GitHub and OpenNeuro repositories.The code and behavioral data are linked to GitHub; raw and processed fMRI data are linked to OpenNeuro.Source: DATA AND CODE AVAILABILITY · p. 10
A.1 STIMULI
The stimuli comprised 18 episodes of This Is Us Season 1, with episodes 2–18 used for training and episode 1 for testing.
A.1 STIMULI
- Episodes 2–18 were training data and episode 1 was test data, with scenes segmented at automatically detected camera-shot changes.Episodes averaged 614 ± 68 scenes; scenes averaged 4.07 ± 4.78 seconds.Source: §A.1 · p. 15
A.2 MOVIE FEATURES
Movie features combined visual, auditory, semantic, and storyline annotations, with normalization and dimensionality reduction applied to embeddings.
A.2 MOVIE FEATURES
- Pretrained CLIP and CLAP models produced 512-dimensional visual and auditory embeddings for each scene.Frame embeddings were averaged across frames; CLAP encoded scene audio.Source: §A.2 · p. 15
- Hierarchical place and character annotations formed 93-dimensional binary semantic vectors used as prediction targets.Categories could encode both a broad label and a finer-grained label, such as indoor and hospital.Source: §A.2 · p. 15
- The four interleaved character storylines supplied the context signal π, with uniform probabilities used when storylines converged or could not be identified.The context was otherwise supplied as a one-hot vector.Source: §A.2 · p. 15
- Embeddings were z-normalized, reduced with principal component analysis, and z-normalized again; context probabilities were normalized to sum to one at each time step.Source: §A.2 · p. 15
A.3 MODEL TRAINING AND CONTEXT MODULATION CONDITIONS
The model receives sequential visual and auditory scene inputs, maintains a recurrent hidden state, and predicts the next scene. Context probabilities can modulate input transformations, recurrent connectivity, or hidden-state output, with recurrent modulation also tested in low-rank form.
A.3 MODEL TRAINING AND CONTEXT MODULATION CONDITIONS
- Training used sequential visual and auditory embeddings from scenes and minimized summed within-episode mean squared error over 100 iterations with Adam and a cosine-decayed learning rate.The learning rate started at 5e-4 and reached 5e-5; one iteration included parameter updates for all 17 episodes.Source: §A.3 · pp. 16–17
- The no-modulation recurrence combines the previous hidden state with a tanh-transformed sum of recurrent and visual and auditory input terms, using λ = 0.7.The model predicts the next scene’s semantic contents from its estimated hidden state.Source: §A.3 · p. 16
- Input modulation replaces the visual and auditory input transformations with four context-specific sets of transformations weighted by context probabilities.Source: §A.3 · p. 16
- Context probabilities weight four context-specific recurrent matrices in full WM modulation, or four unit-rank matrices combined element-wise with a shared matrix in low-rank WM modulation.The low-rank condition replaces the recurrent matrix with the element-wise product of W0 and the weighted sum of context-specific rank-one matrices.Source: §A.3 · p. 16
A.4 MODEL TESTING AND BAYESIAN CONTEXT INFERENCE
At test, the model infers context online while predicting scenes in an unseen episode, using prediction error to update a Bayesian context posterior. The section also specifies how prediction, context inference, and model–brain representation similarity are evaluated.
A.4 MODEL TESTING AND BAYESIAN CONTEXT INFERENCE
- Testing uses unseen episode 1 without a supplied context signal; a recursive Bayesian update carries each posterior forward as the next prior, with a sticky transition prior favoring context persistence.The prior begins uniform across four contexts, and γ is set to 0.9.Source: §A.4 · p. 17
- At each test step, the model compares next-scene predictions under each of four assumed contexts with observed scene semantics, assigning greater likelihood to contexts with lower prediction error.The posterior is computed in log space by combining log prior and log likelihood and normalizing with log-sum-exp.Source: §A.4 · p. 17
A.5 PERFORMANCE
- Next-scene prediction is evaluated by episode loss during training and average Pearson correlation between predicted and observed semantic embeddings at test; context accuracy is the proportion of selected single-storyline scenes assigned the ground-truth context.Correlation values are Fisher z-transformed before averaging and transformed back to r.Source: §A.5 · p. 18
A.6 MODEL–BRAIN REPRESENTATION SIMILARITY
- Model–brain representation similarity compares scene-by-scene representational similarity matrices from model hidden states with voxelwise BOLD activity across 200 cortical regions and 33 participants.The model matrix is averaged across 20 random seeds.Source: §A.6 · p. 18
A.7 FORWARD SIMULATION
Forward simulations trace how hidden states change as a context signal increases from absent to fully specified. The analysis measures trajectory shape, dimensionality, and the separation of context-specific states and recurrent matrices.
A.7 FORWARD SIMULATION
- For fully trained low- and full-rank WM modulation models, hidden-state trajectories were simulated from zero input and state while each context signal increased in 0.001 steps to 1.This produced a 1001-step trajectory for each context.Source: §A.7 · p. 18
- Trajectory tortuosity is the summed stepwise Euclidean distance divided by the start-to-end distance, while dimensionality is assessed from normalized explained variance in the top 20 principal components.The PCA uses concatenated trajectories from the four contexts.Source: §A.7 · p. 18
- Context orthogonality is assessed from pairwise cosine similarity of simulated context endpoints and of the four context-specific recurrent matrices, with the latter tracked across training.Absolute endpoint cosine similarity nearer zero indicates greater orthogonality.Source: §A.7 · pp. 18–19
A.8 CONTEXT-MODULATED EM-RNN
The EM-RNN stores scene representations as keys, values, and context, then combines content- and context-based similarity to retrieve memories. The section also describes comparisons with human retrieval patterns and an alternative that modulates key and query transformations instead.
A.8 CONTEXT-MODULATED EM-RNN
- The EM-RNN stores hidden states as values, input-derived keys as memory addresses, and context signals separately; its test buffer is flushed at episode start and holds up to 1,000 memories.During training, memories are retained across episodes; at test, the buffer contains memories only from episode 1.Source: §A.8 · p. 19
- Retrieval weights combine query-to-key content similarity with current-to-stored-context similarity, while the original EM-RNN uses content similarity alone.The combined weights are softmax-normalized; retrieved memory is a weighted sum of stored values and is blended equally with the current hidden state before prediction.Source: §A.8 · pp. 19–20
A.9 MODEL–HUMAN RETRIEVAL SIMILARITY
- Model–human retrieval similarity is measured by correlating event-level retrieval-weight matrices with human retrieval, after averaging scene-level weights across each pair of 48 events and across scrambled-order groups and seeds.Content-similarity and context-similarity matrices are also correlated separately with the human retrieval matrix.Source: §A.9 · p. 20
A.10 AN ALTERNATIVE MODEL IN WHICH KEY AND QUERY TRANSFORMATIONS ARE MODULATED BY CONTEXT
- An alternative model applies low-rank context modulation to key and query transformations instead of encoding context and using context similarity during memory retrieval.The alternative otherwise retains the original key-query EM system without context modulation.Source: §A.10 · p. 20
B SUPPLEMENTARY FIGURES AND TABLES
The supplementary figures compare context-modulation architectures, memory retrieval selectivity and context congruence, and softmax settings; the tables describe comparison measures.
Figure S1
- Figure S1 contrasts an unmodulated GRU EM-RNN with full WM modulation that sums four recurrent matrices weighted by context probability.The unmodulated model predicts visual embeddings from visual scene embeddings and retrieves through its key-value system without explicitly representing context in the EM buffer.Source: Figure S1 · p. 21
Figure S2
- Across significant cortical regions, model–brain RSM similarity is higher with low-rank than full WM modulation.The comparison uses paired tests across 33 participants and reports FDR-corrected regional t statistics.Source: Figure S2 · p. 22
Figure S3
- WM+EM modulation yields more selective retrieval and retrieves memories with greater context similarity than WM modulation alone.It has fewer retrieval-matrix values in the 1e-5 to 1e-1 bins and a relatively higher proportion in the [1e-1, 1] bin, though that proportion remains smaller than under WM modulation.Source: Figure S3 · p. 22
Figure S4
- Figure S4 varies softmax temperature around the main-analysis value of τ = 0.1, with larger τ reducing retrieval selectivity and no softmax producing nonselective retrieval.The no-softmax result is shown in the right panel.Source: Figure S4 · p. 23
Figure S5; Tables S1–S2
- Figure S5 tracks test-time Bayesian context-inference accuracy across training iterations, while Tables S1 and S2 define the measures used to compare modulation conditions.Table S1 specifies next-scene prediction, context-inference, and model–brain representation measures after 100 iterations; Table S2 also defines model–human retrieval similarity and describes WM and EM modulation.Source: Figure S5 · p. 23, Table S1 · p. 23, Table S2 · p. 24