Memory-Augmented Brain Encoding — Teaching an fMRI Model to Remember
The core idea: a content-based memory module that retrieves relevant past context and feeds it back into a state-of-the-art brain-encoding model.
For the last few months I've been working on a research project at Stockton University, and I finally want to write it down plainly — not as a paper abstract, but as the story of what I actually built and why. The project is called Memory-Augmented Brain Encoding, and the whole thing lives on GitHub at Mrabbi3/memory-augmented-brain-encoding.
Below the coffee-chat version, I've included the entire manuscript — every section, every figure, every table — so if you want to read the real thing end to end, you can do it right here. You can also grab the original:
The problem, in one breath
An encoding model predicts brain activity from whatever a person is experiencing. Show someone a movie inside an fMRI scanner, record which parts of their cortex light up, and train a model to go from the stimulus → the brain response. If the model is good, it has learned something real about how the brain represents the world.
The current state of the art is a model called TRIBE v2 — the model that won 1st place out of 263 teams in the Algonauts 2025 challenge. It fuses video (CLIP), audio (Whisper), and text (LLaMA) and predicts activity across the whole cortical surface. It's genuinely impressive. But it has one blind spot: it treats every moment of the movie independently. It maps this instant's features to this instant's brain response, with no memory of the ten minutes of story that came before.
That's neurobiologically implausible. Whether a scene feels tense, funny, or heartbreaking depends on plot you accumulated over the previous several minutes — and we know from neuroscience that higher-order cortex integrates information over long timescales, while sensory cortex lives in the moment. This is the temporal receptive window (TRW) hierarchy.
So the whole project comes down to one focused question:
Does adding learned, content-based memory over past context improve encoding — and where in the cortex does it help most?
The short answer: memory doesn't help everywhere, and that's exactly the point. It selectively improves the association networks that theory says integrate over long timescales — Frontoparietal Control (ΔR = +0.048, p < 10⁻⁵) and Salience/Ventral Attention (ΔR = +0.044, p = 0.0003) — while leaving (or slightly hurting) primary sensory cortex. That dissociation, not a bump in the overall average, is the real finding.
The full paper
What follows is the complete manuscript — "Memory-Augmented Brain Encoding: Improving Neural Prediction through Temporal Context Integration" — rebuilt for the web with the actual figures and results. The original PDF is here.
Abstract
Brain encoding models predict neural responses from stimulus features, but current approaches process each time point independently, ignoring the temporal context that shapes how the brain integrates information during naturalistic stimulation. I propose a memory-augmented encoding architecture that extends state-of-the-art brain encoding models with a cross-attention memory module, enabling the model to integrate past stimulus representations when predicting current neural activity. Using the Algonauts 2025 dataset — fMRI recordings of subjects watching Friends episodes, parcellated into 1,000 Schaefer regions — I show that memory augmentation selectively improves encoding in higher-order association networks. Specifically, Frontoparietal Control (ΔR = +0.048, p < 10⁻⁵) and Salience/Ventral Attention (ΔR = +0.044, p = 0.0003) show the strongest memory benefit, while primary sensory regions (Visual, Somatomotor) do not benefit. This network-specific dissociation aligns with evidence that higher-order networks operate over longer temporal integration windows during narrative comprehension. Systematic ablations across five design dimensions identify an optimal compact configuration — 8 memory slots, 2 attention heads, learned gating — achieving maximal benefit with only 42% parameter overhead (2.7M vs. 1.9M baseline). Cross-task evaluation on Movie10 data confirms the benefit generalizes beyond the training stimulus (ΔR = +0.020, 57.1% of parcels improved). The result is both an engineering contribution — a simple, effective memory module for encoding models — and a neuroscientific insight connecting temporal context integration to the cortical network hierarchy.
I. Introduction
Brain encoding models map stimulus features to predicted neural responses. Recent advances — exemplified by TRIBE v2 winning the Algonauts 2025 Challenge — achieve remarkable accuracy using multimodal feature extractors (CLIP ViT-L/14 for vision, 768-dim; Whisper large-v2 for audio, 1,280-dim; LLaMA 3.2 for language, 2,048-dim) combined with a learned mapping to cortical-surface predictions across 20,484 vertices.
But a fundamental limitation persists: current encoding models treat each repetition time (TR) independently. At each moment, the model maps the current stimulus features to the predicted brain response without considering preceding context. The brain does the opposite — it continuously integrates information over time, maintaining representations of past events that shape current processing. This matters most during naturalistic stimulation like watching television, where comprehension requires tracking characters, plot threads, and causal relationships across extended durations.
The consequences of the limitation are visible in exactly the regions you'd predict. When you decompose TRIBE v2's own learned weights, the spatial structure recovers the familiar functional networks:
Figure 1 — ICA decomposition of TRIBE v2's learned weight matrices. Each independent component
(IC1–IC5) isolates a spatial pattern the pretrained model relies on: auditory, language, motion,
default-mode, and visual systems. Reproducing these components confirms the pretrained model
encodes recognizable functional networks — the substrate our memory module plugs into.
And when you project the strength of those learned weights onto the cortex, the weakest encoding sits precisely in the higher-order regions:
Figure 2 — TRIBE v2 weight norms on the cortical surface (darker = stronger learned mapping,
lighter = weaker). The inferior frontal cortex (weight norm = 1.0008) and the multisensory
temporoparietal junction (1.0278) rank lowest — the regions most involved in narrative
comprehension and temporal integration. This motivates the central hypothesis: missing temporal
context, not insufficient model capacity, is what limits encoding in these areas.
My proposal is a simple architectural extension: a memory-augmented encoding module that maintains a circular buffer of past stimulus representations and uses multi-head cross-attention to retrieve relevant context when predicting current activity. A learned gate, initialized low, lets the model gradually discover how much temporal context is useful. It adds minimal parameters while providing the temporal-integration capacity current models lack.
Contributions. (1) A memory-augmented encoding architecture with cross-attention retrieval and learned gating, applicable to any existing brain encoding model. (2) Empirical evidence that memory selectively improves higher-order association networks (Frontoparietal, Salience) while leaving primary sensory areas unaffected — a network-specific dissociation. (3) Systematic ablations across five design dimensions identifying the optimal, minimal configuration. (4) Cross-task generalization showing the benefit transfers from Friends episodes to Movie10 stimuli.
II. Related work
Brain encoding models. Early work used hand-crafted features and ridge regression; more recent models use deep-network representations as the feature space — CNN layers for vision, language-model activations for text. TRIBE v2 is the current state of the art, using multimodal features with a learned linear mapping to the cortical surface. Critically, all of these treat each time point independently.
Temporal processing in the brain. The importance of temporal context is well established. Event-segmentation theory holds that the brain carves continuous experience into discrete events, with hippocampal and prefrontal regions tracking narrative structure. Hasson et al. showed that temporal receptive windows vary systematically across the cortical hierarchy — primary sensory areas integrate over hundreds of milliseconds, while default-mode and frontoparietal regions integrate over tens of seconds to minutes. Simony et al. showed that inter-subject neural coupling during narrative comprehension is strongest in regions with long TRWs. Despite this, most encoding models still process time points independently.
Memory-augmented networks. Memory augmentation has proven effective across domains — Neural Turing Machines and Differentiable Neural Computers introduced external memory with attention-based read/write; transformers implicitly maintain context through self-attention. My approach applies these ideas to brain encoding, where temporal structure has direct neurobiological meaning. The key difference from a standard transformer is a compact circular buffer with cross-attention rather than full self-attention — cheap enough for sequential fMRI prediction.
III. Methods
A. Dataset
Everything runs on the Algonauts 2025 dataset: fMRI from subjects watching Friends episodes and Movie10 clips, parcellated into 1,000 regions with the Schaefer atlas and the 7-network Yeo parcellation. Each parcel is assigned to one of seven canonical networks:
| Yeo network | Parcels |
|---|---|
| Visual | 162 |
| Somatomotor | 194 |
| Dorsal Attention | 122 |
| Salience / Ventral Attention | 121 |
| Limbic | 60 |
| Frontoparietal Control | 128 |
| Default Mode | 212 |
The main analyses use Subject 1: 478 TRs for Friends (~12 minutes, TR = 1.49 s) and 399 TRs for Movie10. Here's what the raw signal actually looks like:
Figure 3 — Real fMRI (Subject 1 watching Friends, ses-001, 478 TRs, 1,000 parcels). Top: z-scored
BOLD time courses for four representative parcels — noisy, but clearly structured. Bottom:
inter-parcel correlation matrix for the first 100 parcels, showing the block structure and
temporal dependencies that a memory module can exploit.
B. Feature extraction
Stimulus features come from TRIBE v2's pretrained encoders (CLIP ViT-L/14, 768-dim; Whisper large-v2, 1,280-dim; LLaMA 3.2 1B, 2,048-dim), aligned to TR resolution and concatenated. I apply a hemodynamic response function (HRF) shift of 4 TRs to account for neurovascular delay.
For experiments where end-to-end extraction wasn't feasible, I use Algonauts' pre-extracted features (3,584-dim audio/visual) plus enhanced surrogate features built from five independent temporal sources derived from the fMRI signal itself — PCA temporal patterns (50-dim), ICA components (30-dim), temporal derivatives (50-dim), a slow narrative arc via Gaussian smoothing (50-dim), and amplitude-envelope modulation (20-dim).
C. Memory-augmented encoding architecture
Figure 4 — The full architecture. Multimodal features (CLIP + Whisper + LLaMA, 4096-d
concatenated) are projected to a 256-d embedding. The memory module keeps a circular buffer of 8
past embeddings; the current embedding queries that buffer through 2-head cross-attention, and a
learned gate σ(α) ≈ 0.12 controls how much retrieved memory is mixed back in via a gated residual.
A 3-layer cortical decoder maps the result to 1,000 Schaefer parcels.
The model has three stages:
1) Feature projector. A two-layer MLP with LayerNorm and GELU projects the concatenated features fₜ into a common embedding space:
eₜ = Proj(fₜ) ∈ ℝ^d (d = 256 by default)
2) Memory module. A circular buffer M ∈ ℝ^{M×d} of size M (default M = 8) stores past embeddings. At each step t, the current embedding queries the buffer via multi-head cross-attention (H = 2 heads):
aₜ = MultiHead(eₜ, M₁:ₖ, M₁:ₖ) k = min(t, M) valid entries
A learnable scalar gate α controls the memory contribution through a gated residual:
e′ₜ = LayerNorm( eₜ + σ(α) · aₜ ) α initialized to −2 → σ(−2) ≈ 0.12
A feed-forward block with a residual connection produces the final embedding
zₜ = LayerNorm(e′ₜ + FFN(e′ₜ)). After each step, eₜ is written to the buffer at position
t mod M (detached from the computation graph, so memory carries information without
backpropagating through all of history).
3) Cortical decoder. A three-layer MLP maps the memory-augmented embedding to predicted
activity for all P = 1,000 parcels: ŷₜ = Dec(zₜ) ∈ ℝ^P.
D. Training
Models train with MSE loss and AdamW (lr = 3×10⁻⁴, weight decay 0.01) with cosine annealing. TRs are processed sequentially within contiguous windows of 100 TRs (stride 50) to maintain memory state. Gradients are clipped at norm 1.0. Data is split temporally — 80% train / 20% validation, no shuffling. Final models train for 50 epochs; ablation studies for 20.
E. Evaluation metric
The primary metric is per-parcel Pearson correlation rₚ between predicted and actual fMRI. The memory benefit is the change from adding memory to an otherwise identical baseline:
ΔRₚ = rₚ(memory) − rₚ(baseline)
I report mean ΔR across parcels and network-level statistics using one-sample t-tests (H₀: ΔR = 0, one-sided). A caveat worth stating up front: absolute fMRI encoding correlations are modest by nature — predicting brain activity from a movie is hard and nobody gets high absolute numbers — so the honest headline is ΔR, not raw accuracy.
IV. Experimental results
A. Pipeline validation on synthetic data
Before trusting anything on real brains, I validated the whole pipeline on synthetic fMRI with known temporal dependencies: 300 TRs with five narrative events, exponential memory decay in the encoding weights, and explicit "narrative callback" effects in designated hippocampal vertices, where later events reference earlier ones.
Figure 5 — Pipeline validation on synthetic data. The memory model (blue) reaches R = 0.086 vs.
the baseline's R = 0.054, with the biggest gains in the "cortical" units that carry long-range
callbacks (ΔR_cortex = +0.056) and smaller gains in hippocampal units (ΔR_hippo = +0.008). The
learned gate opens during training and settles near 0.12 — the model discovers, on its own, that
past context is useful. This is the sanity check: when temporal dependencies are real and known,
memory finds them.
B. Memory consistently improves real brain encoding
On real data, across multiple feature conditions, memory augmentation consistently improves encoding:
Figure 6 — First encoding results on real fMRI (Friends, Subject 1, surrogate features). Top:
training curves, encoding accuracy, and the gate opening. Middle: per-parcel memory benefit —
614 of 1,000 parcels show ΔR > 0. Bottom: the ΔR distribution (mean = +0.026), the top-20
most-improved parcels, and summary statistics.
The benefit isn't an artifact of the surrogate features — it holds with the real pre-extracted audio/visual features too:
Figure 7 — Encoding with real pre-extracted features (3,584-dim). The consistency between
surrogate features (ΔR = +0.026) and real features (ΔR = +0.021) confirms the memory benefit is
not a construction artifact; 574 of 1,000 parcels improved.
| Condition | Features | R (mem) | R (base) | ΔR | Gate |
|---|---|---|---|---|---|
| Friends (surrogate) | 1024-d | 0.037 | 0.011 | +0.026 | 0.119 |
| Friends (real) | 3584-d | 0.023 | 0.001 | +0.021 | 0.118 |
| Movie10 (surrogate) | 2048-d | 0.130 | 0.110 | +0.020 | 0.124 |
| Average | +0.022 | 0.120 |
Table I — Encoding results across conditions. The learned gate converges to σ(α) ≈ 0.12 everywhere, suggesting ~12% memory contribution is optimal.
C. The headline: network-specific memory benefits
This is the finding I care about most. The benefit is not uniform across cortex — it dissociates cleanly by functional network:
Figure 8 — Network-level analysis. Top: mean ΔR per Yeo network with standard-error bars and
significance markers (*** p < 0.001, ** p < 0.01, * p < 0.05). The gradient runs from
Frontoparietal / Salience (strong benefit) down to Visual / Somatomotor — mirroring the cortical
temporal-receptive-window hierarchy. Middle: violin plots of the full per-parcel ΔR distributions.
Bottom: a small right-vs-left hemisphere difference (not significant, t = −1.30, p = 0.195) and a
per-parcel Memory-vs-No-Memory scatter colored by network.
| Network | N | ΔR | p | Sig. |
|---|---|---|---|---|
| Frontoparietal Control | 128 | +0.048 | < 10⁻⁵ | *** |
| Salience / Ventral Attention | 121 | +0.044 | 0.0003 | *** |
| Limbic | 60 | +0.013 | 0.187 | |
| Default Mode | 212 | +0.003 | 0.372 | |
| Dorsal Attention | 122 | −0.016 | 0.072 | |
| Somatomotor | 194 | −0.025 | 0.001 | *** |
| Visual | 162 | −0.038 | < 10⁻⁴ | *** |
Table II — Memory benefit by Yeo network (optimal model, Friends), sorted by ΔR. Higher-order association networks improve significantly; primary sensory networks are significantly impaired.
Why does this make neuroscientific sense? Frontoparietal and Salience networks integrate information over longer windows during narrative comprehension — they maintain representations of task context, character goals, and plot structure. That's exactly the kind of information a memory buffer supplies. Primary sensory regions, by contrast, respond to immediate physical properties already captured by the current-moment features — so memory is not just unnecessary there, it adds noise. The dissociation replicates with the real pre-extracted features too:
Figure 10 — The network-level pattern replicated with real pre-extracted features. Somatomotor
(*** p < 10⁻⁸), Visual (*** p = 0.0002), and Dorsal Attention (** p = 0.002) show
significant effects, confirming the dissociation is robust across feature types.
D. Where it lives on the cortex
Projected onto the cortical surface, the spatial pattern confirms the network story:
Figure 9 — Memory benefit (ΔR) on the cortical surface (fsaverage), four views. Red = memory
improves encoding; blue = memory impairs it (scale ±0.19). Red clusters concentrate in lateral
and medial prefrontal cortex, lateral parietal cortex, and the temporoparietal junction — the
Frontoparietal and Salience territory. Blue clusters localize to occipital visual cortex and
pericentral somatomotor cortex. The map is the hypothesis, drawn on a brain.
E. Ablation studies — what actually matters
I swept five design dimensions, one factor at a time, holding the rest at defaults:
Figure 11 — Ablations across five design choices. Each panel plots encoding correlation (blue) and
memory benefit ΔR (green), with the optimum boxed in orange. The clear winners: hidden dim 256,
2 attention heads, learned gating, and a small 8-slot memory. Bigger is not better — larger hidden
dims and larger buffers hurt.
| Dimension | Best | ΔR | % Improved |
|---|---|---|---|
| Hidden dimension | 256 | +0.033 | 62.7% |
| Attention heads | 2 | +0.012 | 54.8% |
| Gate mechanism | Learned | +0.011 | 54.8% |
| Memory size | 8 slots | +0.004 | 51.3% |
| Sequence length | 100 TRs | −0.007 | 49.2% |
Table III — Ablation results (20 epochs each). Each row varies one dimension.
The single most impactful choice is hidden dimension: 256 gives ΔR = +0.033, while 512 collapses to +0.006 and 768 goes negative (−0.014). With only 478 TRs of training data, compact models generalize; big ones overfit. Two heads beat one (−0.005) and beat four or eight (+0.004 each). A learned gate beats fixed-0.5, fixed-1.0, and no-gate. And small memory beats large — 8 slots (+0.004) beats 128 slots (−0.036), because distant buffer entries just inject noise.
Figure 12 — Left/center: the learned gate converges to σ(α) ≈ 0.119 regardless of memory size or
head count — a strikingly stable "88:12" current-to-memory ratio. Right: parameter efficiency —
hidden dim 256 (upper-left) achieves the best ΔR with the fewest parameters, while larger models
drift down and to the right. The optimal model is 2.7M parameters, only 42% above the 1.9M
baseline.
That 88:12 ratio is oddly consistent. Across every configuration, the model settles on mixing in about 12% memory. It may reflect the actual balance of bottom-up vs. top-down information during passive viewing — a nice, testable little hypothesis for future work.
F. Cross-task generalization
Does the benefit transfer to a different stimulus? Partly — and the way it fails is informative:
| Condition | R (mem) | R (base) | ΔR |
|---|---|---|---|
| Zero-shot (Friends → Movie10) | −0.019 | −0.013 | −0.006 |
| Trained on Movie10 | +0.130 | +0.110 | +0.020 |
Table IV — Cross-task transfer. Zero-shot transfer of a Friends-trained model does not help (the memory learns stimulus-specific temporal patterns), but a fresh model trained on Movie10 with the optimal config recovers a clear benefit (ΔR = +0.020, 57.1% improved) — closely matching Friends. Temporal-context integration is a general benefit; the specific learned patterns are not.
G. Progression, and an honest caveat
Figure 13 — Final results with the optimal configuration (hidden 256, memory 8, heads 2, learned
gate; 2,682,601 params). Top: encoding accuracy, sorted per-parcel benefit (495/1,000 improved),
and the flat, stable gate. Middle: the memory benefit across development days, the Movie10 transfer,
and the ΔR distribution. Bottom: the network-level dissociation, still strong — Frontoparietal
(*** p < 10⁻⁵) and Salience (*** p = 0.0003) on top, Visual and Somatomotor negative.
Here's the caveat I want to be transparent about. At this fully-optimized config trained for 50 epochs on Friends, the overall average ΔR is essentially zero (−0.0002, 49.5% of parcels improved) — the baseline caught up on the global mean. But look at the bottom panel: the network-specific dissociation is undiminished. Memory still significantly helps Frontoparietal (+0.048) and Salience (+0.044) and still significantly hurts Visual and Somatomotor. In other words, a single global average is the wrong summary statistic for this effect — the benefit is real and structured, it's just localized. That's why the paper leads with the network breakdown, not the grand mean.
V. Discussion
Network dissociation and the temporal hierarchy. The split between Frontoparietal/Salience (benefit) and Visual/Somatomotor (no benefit) aligns with the TRW hierarchy. Higher-order regions maintain representations over seconds to minutes, integrating narrative elements and tracking intentions; sensory regions respond to immediate properties already captured by current features. Rather than measuring neural timescales directly, this is a model-based characterization: a temporal-integration component selectively improves prediction exactly where TRWs are long.
The 88:12 ratio. The gate's convergence to ≈ 0.12 across every experiment is noteworthy. The
optimal balance between current (88%) and past (12%) information may reflect the real ratio of
bottom-up to top-down processing during passive viewing. Whether it shifts with task demands or
stimulus complexity is an open question.
Implications for encoding-model design. The results suggest a network-specific strategy: apply memory augmentation only to Frontoparietal and Salience parcels, and keep memoryless encoding for Visual and Somatomotor. A targeted architecture like that could beat the uniform one tested here.
Connection to AI memory architectures. Different AI memory strategies — sliding window, summarization, retrieval-augmented generation (RAG), hybrid — may map differentially onto brain networks' temporal profiles. Visual cortex may best match a sliding window; Default Mode may match summarization; Frontoparietal may match RAG. I'm pursuing this systematic cross-strategy comparison in a BioMemBench framework, connecting AI agent memory design to biological memory systems.
Limitations. Training data is limited (478 TRs), constraining capacity. The Schaefer parcellation includes some subcortical structures. The real pre-extracted features capture only audio/visual modalities. I evaluate on a single subject. And the circular buffer is deliberately simple — hierarchical or consolidation-based memory may better capture multi-timescale processing.
VI. Conclusion
I presented a memory-augmented brain encoding architecture that selectively improves neural prediction in higher-order networks (Frontoparietal ΔR = +0.048, p < 10⁻⁵; Salience ΔR = +0.044, p = 0.0003) while leaving primary sensory areas unaffected. This dissociation aligns with the cortical temporal-receptive-window hierarchy and offers a model-based way to characterize how brain networks integrate information over time. The optimal configuration — 8 slots, 2 heads, learned gate, 2.7M parameters — is compact and drops into any existing encoding model. Cross-task generalization to Movie10 confirms the effect is not stimulus-specific.
The core result already feels solid to me: you can teach a state-of-the-art brain encoder to remember, and it improves prediction in exactly the parts of cortex where theory says it should, with only a small parameter cost. For a project that began as one focused question — does memory help, and where? — that's a satisfying answer.
Selected references
- d'Ascoli, Bhatt, et al. TRIBE v2: Multimodal brain encoding model. Meta FAIR, 2025.
- Hasson, Yang, Vallines, Heeger, Rubin. A hierarchy of temporal receptive windows in human cortex. J. Neurosci., 2008.
- Baldassano, Chen, Zadbood, Pillow, Hasson, Norman. Discovering event structure in continuous narrative perception and memory. Neuron, 2017.
- Simony, Honey, Chen, et al. Dynamic reconfiguration of the default mode network during narrative comprehension. Nat. Commun., 2016.
- Schaefer et al. Local-global parcellation of the human cerebral cortex. Cereb. Cortex, 2018.
- Yeo et al. The organization of the human cerebral cortex estimated by intrinsic functional connectivity. J. Neurophysiol., 2011.
- Graves, Wayne, Danihelka. Neural Turing Machines. arXiv:1410.5401, 2014.
- Vaswani et al. Attention is all you need. NeurIPS, 2017.
- The Algonauts Project 2025. algonautsproject.com
Author: MD Rabbi, Department of Computer Science, Stockton University (rabbim@go.stockton.edu). Manuscript prepared April 2026.
📄 Download the full paper (PDF) · 💻 Code, notebooks & figures on GitHub
Related: Ai Engineering · Software Engineering
