Predictive Associative Memory: Retrieval Beyond Similarity Through Temporal Co-occurrence

AI summary

Overview

• Research area: Machine learning / neural memory architectures — specifically associative memory and retrieval, positioned against retrieval-augmented generation (RAG), Memory Networks, Hopfield networks, knowledge graphs, the successor representation, and spreading activation.
• Technical level: Intermediate. The paper's argument is stated in largely plain language (stairs and slipping, sunscreen and holidays), but the architecture assumes familiarity with embedding spaces, JEPA-style latent prediction, contrastive/InfoNCE training, and AUC-style evaluation.
• Scope (one sentence): The paper proposes Predictive Associative Memory (PAM), an architecture in which a JEPA-style "Inward" predictor learns memory retrieval from temporal co-occurrence rather than representational similarity, and it evaluates that predictor in isolation as an episodic recall system on a synthetic benchmark.

What This Paper Is About

Most neural memory systems retrieve by similarity: given a query, they return whatever stored item sits closest in embedding space. The paper argues this misses a defining property of biological memory, where things become linked simply by having been experienced close together in time — stairs do not resemble a slip, yet one evokes the other. The goal is to build a mechanism that learns association structure directly from temporal co-occurrence in an experience stream, and to test whether it can recall associations that similarity search cannot.

Key Contributions

1. The Inward JEPA framing. A formalisation of memory retrieval as a predictor that maps a current composite state to the region of embedding space containing temporally co-occurring past states, positioned as the complement of the standard "Outward" JEPA that predicts future sensory states. Both share an encoder and a prediction-in-latent-space mechanism but target different spaces.
2. An architecture for associative memory without explicit structure construction. Temporal co-occurrence within a window τ is used as the sole training signal, requiring no hand-built graph, no entity extraction pipeline, and no downstream task supervision — explicitly contrasted with knowledge graphs and Memory Networks.
3. A recall-versus-retrieval evaluation paradigm. The paper argues associative memory should be judged on faithfulness of recall for experienced associations, not generalisation to unseen ones, and defines metrics for the corresponding failure modes: false association, loss of specificity, and signal corruption.
4. A synthetic benchmark and empirical validation of the Inward channel in isolation, including controls (temporal shuffle, specificity against similar-but-not-associated distractors, held-out query states) and a reported negative result on cross-episode creative bridging.

Main Findings

• Top retrieval is usually a true temporal associate. The predictor's top retrieval is a true temporal associate 97% of the time, reported as Association Precision@1 = 0.970.
• Cross-boundary recall works where similarity fails. The predictor achieves cross-boundary Recall@20 = 0.421 on pairs from different rooms, where cosine similarity scores zero.
• Discrimination between experienced-together and never-experienced-together states. Discrimination AUC = 0.916 for the predictor versus 0.789 for cosine similarity.
• The advantage survives when similarity is uninformative. Restricted to cross-room pairs, the predictor achieves AUC = 0.849 while cosine similarity is 0.503, described as chance.
• Specificity against distractors is confirmed. Specificity controls testing true temporal associates against similar-but-not-associated distractors give AUC = 0.848 versus cosine 0.732.
• The signal is genuinely temporal, not geometric. A temporal shuffle control — shuffling temporal order within trajectories — collapses cross-boundary recall by 90%, replicated across training seeds.
• Recall is anchor-specific. A held-out query-state evaluation shows the predictor recalls associations from experienced viewpoints but not from novel ones, which the paper describes as consistent with the episodic character of the memory system.
• Results are stable across seeds and query selections. Standard deviations are below 0.006 across training seeds and at most 0.012 across query selections.
• A linear compatibility function is not sufficient. The bilinear baseline, trained on the same temporal co-occurrence data with the same InfoNCE objective, reaches above-chance overall discrimination (AUC = 0.791) but drops to chance on cross-room pairs (AUC = 0.514) and scores zero on cross-boundary recall and specificity. The paper concludes the nonlinear predictor is necessary specifically for associations that cross representational boundaries.
• Creative bridging is a negative result. The proposed transitive bridging mechanism (trajectory 1 visits A then B, trajectory 2 visits B then C, so from A the predictor could reach C) requires entity persistence across episodes, which the paper states the benchmark does not yet provide.

Methodology in Plain English

The researchers built a synthetic world of 20 rooms in a 128-dimensional embedding space, with each room defined by a cluster centroid at scale 2.0 and rooms well separated so that states in the same room are similar and states in different rooms are distant. Fifty objects are distributed across the rooms (5 per room, with some objects shared between rooms), and objects contribute features at scale 1.5 that modulate the room embedding.

Experience is generated as 500 trajectories of 100 timesteps each — 50,000 total states. At each timestep the agent occupies a room and perceives the objects present, producing a composite state embedding. Trajectories cross room boundaries, which creates temporal co-occurrences between states from different rooms. Associations are defined by a co-occurrence window of τ = 5 timesteps, yielding 242,264 total associations.

The predictor itself is a 4-layer MLP with residual connections (128 → 1024 → 1024 → 1024 → 128), GELU activations, residual connections in layers 2–3, and layer normalisation on the output, totalling 2.36M parameters. The authors note a transformer was not used because the predictor operates on single composite states rather than sequences. It is trained with InfoNCE loss using in-batch negatives (batch size 512, giving 511 negatives per positive), a cosine learning rate schedule from 5 × 10⁻⁴ to 1 × 10⁻⁵, and temperature annealing from 0.15 to 0.05 over 500 epochs.

Two baselines are used: standard k-nearest-neighbour retrieval by cosine distance (no training, representing the similarity-only approach used by RAG and dense retrieval), and a learned bilinear form trained on the same data with the same objective, included as a test of whether the form of the mapping (linear vs nonlinear) matters rather than as a capacity-matched competitor.

Two points about the setup are emphasised. First, the predictor is trained on all temporal associations in the experience stream, consistent with the recall paradigm — the system is trained on its complete experience and tested on whether it faithfully recalls it. Second, the predictor is evaluated over a fixed synthetic embedding space rather than inside a full JEPA with joint encoder training; since there is no raw data for an encoder to process, the stop-gradient and EMA mechanisms are not required here, though the paper says they become necessary when the predictor is integrated with a co-evolving encoder.

The architecture also specifies two mechanisms not exercised in the benchmark: familiarity normalisation, where association strength is the residual above a running-average co-occurrence baseline (described as sensory adaptation applied to memory formation), and adaptive decay, where the forgetting rate increases with average association density to maintain retrieval precision.

Why This Matters

Impact on research. The paper reframes memory retrieval as traversal of learned associative structure rather than search by similarity, and supplies the differentiable, trainable mechanism that the spreading activation literature (Collins and Loftus, 1975) described but never implemented as a learned operation. It also draws a clean evaluation distinction between retrieval generalisation and episodic recall, arguing that strong performance on unexperienced associations would indicate false-link formation rather than success. The dual-channel picture — Outward for similarity, Inward for association — yields an architectural prediction: episodic specificity emerges from the intersection of the two channels, since similarity alone returns all drills and association alone returns everything from Tuesday, while their intersection returns that specific drill on that specific occasion. The paper is explicit that the present work validates only the Inward channel in isolation.

Real-world applications (as framed by the paper's argument):

• Memory for embodied agents that accumulate experience and need associations reflecting the structure of that experience rather than the structure of a particular downstream task.
• Retrieval systems where useful memories are distant in embedding space from the query — the paper's examples are cross-modal and cross-context cases that similarity retrieval cannot recover.
• Personalised episodic recall over a continuous experience stream, where the goal is faithful reactivation of specific past conjunctions rather than generalisation.
• A memory substrate for planning systems: the paper notes an active inference agent could use PAM-style associative retrieval as the memory over which it plans.

Industry relevance. The direct point of contrast is deployed similarity-based retrieval, including RAG and dense passage retrieval. The paper argues the limitation is architectural — similarity retrieval requires the very proximity that is absent in exactly the cases where biological memory excels — and that query transformation methods such as hypothetical document embeddings and query rewriting leave the retrieval mechanism itself similarity-based, requiring the transformation to guess the right direction in advance. The retrieval formulation is also designed to be operationally cheap: a single forward pass through the predictor followed by a nearest-neighbour lookup in the predicted region, which the paper notes scales to large memory banks via approximate nearest-neighbour indices, with the predictor simply redirecting which point the index is queried from. The paper also cites recent survey work reporting that reasoning failures in large language models are particularly acute in embodied and causal domains where grounded experience is absent (Song et al., 2026).

Future Directions

1. Integrating the Inward predictor with a learned encoder. The benchmark uses a fixed synthetic embedding space; the paper identifies integration with a co-evolving encoder (for example, V-JEPA over video) as the natural next step, at which point the stop-gradient and EMA mechanisms become necessary to prevent representational collapse.
2. Achieving creative bridging. The transitive recombination result was negative because it requires entity persistence across episodes, which the benchmark does not provide. Making multi-hop retrieval bridge separate trajectories remains an open problem.
3. Salience-weighted co-occurrence. The paper isolates temporal co-occurrence as the foundational signal and proposes that affective and arousal signals — predicted to emerge from the architecture's own dynamics, such as prediction error magnitude, novelty relative to the world model, and goal-state discrepancy — would enter as additional dimensions of the composite state.
4. Richer uncertainty representation in the predictor. The current predictor outputs a single point (the predicted centroid of the association neighbourhood), with the "region" defined implicitly by a retrieval threshold. Predicting both a mean and variance as a Gaussian in embedding space is proposed as a natural extension.

Target Audience

Researchers and engineers working on memory architectures for neural systems — particularly those building retrieval systems, embodied agents, or world models — will get the most from this paper, since its central claim is directed at the similarity-retrieval assumption underlying RAG and dense retrieval. It is also relevant to cognitive scientists and neuroscientists interested in hippocampal indexing theory, complementary learning systems, and spreading activation, because it offers a computational implementation of associative structure learning that those models described but did not mechanise. Readers will need enough background to follow embedding-space geometry, JEPA-style latent prediction, and contrastive training; the argument itself is presented without heavy jargon.

Authors’ abstract

Current approaches to memory in neural systems rely on similarity-based retrieval: given a query, find the most representationally similar stored state. This assumption -- that useful memories are similar memories -- fails to capture a fundamental property of biological memory: association through temporal co-occurrence. We propose Predictive Associative Memory (PAM), an architecture in which a JEPA-style predictor, trained on temporal co-occurrence within a continuous experience stream, learns to navigate the associative structure of an embedding space. We introduce an Inward JEPA that operates over stored experience (predicting associatively reachable past states) as the complement to the standard Outward JEPA that operates over incoming sensory data (predicting future states). We evaluate PAM as an associative recall system -- testing faithfulness of recall for experienced associations -- rather than as a retrieval system evaluated on generalisation to unseen associations. On a synthetic benchmark, the predictor's top retrieval is a true temporal associate 97% of the time (Association Precision@1 = 0.970); it achieves cross-boundary Recall@20 = 0.421 where cosine similarity scores zero; and it separates experienced-together from never-experienced-together states with a discrimination AUC of 0.916 (cosine: 0.789). Even restricted to cross-room pairs where embedding similarity is uninformative, the predictor achieves AUC = 0.849 (cosine: 0.503, chance). A temporal shuffle control confirms the signal is genuine temporal co-occurrence structure, not embedding geometry: shuffling collapses cross-boundary recall by 90%, replicated across training seeds. All results are stable across seeds (SD < 0.006) and query selections (SD $\leq$ 0.012).