Human-inspired Episodic Memory for Infinite Context LLMs

Authors: Zafeirios Fountas, Martin Benfeghoul, Adnan Oomerjee, Fenia Christopoulou, Gerasimos Lampouras, Haitham Bou Ammar, Jun Wang
Primary area: foundation or frontier models, including LLMs
Keywords: large language models, long context, retrieval, episodic memory, event cognition, training-free

Abstract:
Large language models (LLMs) have shown remarkable capabilities, but still struggle with processing extensive contexts, limiting their ability to maintain coherence and accuracy over long sequences. In contrast, the human brain excels at organising and retrieving episodic experiences across vast temporal scales, spanning a lifetime. In this work, we introduce EM-LLM, a novel approach that integrates key aspects of human episodic memory and event cognition into LLMs with no fine-tuning, enabling them to handle practically infinite context lengths while maintaining computational efficiency. EM-LLM organises sequences of tokens into coherent episodic events using a combination of Bayesian surprise and graph-theoretic boundary refinement in an online fashion. When needed, these events are retrieved through a two-stage memory process, combining similarity-based and temporally contiguous retrieval for efficient, human-inspired access to relevant information. Experiments on the LongBench and $\infty$-Bench benchmarks demonstrate EM-LLM's superior performance, consistently outperforming the state-of-the-art retrieval model InfLLM across various baseline LLMs. In addition, EM-LLM outperforms its popular counterpart, RAG, in a wide range of tasks, while requiring similar resources. Notably, EM-LLM's performance even surpasses full-context models in most tasks, while successfully performing retrieval across 10 million tokens -- a scale computationally infeasible for such models. Finally, our analysis reveals strong correlations between EM-LLM's event segmentation and human-perceived events, suggesting parallels between this artificial system and its biological counterpart, thereby offering a novel computational framework for exploring human memory mechanisms.

Decision: accept (poster)
Reviewer scores: 5, 8, 5, 5

Reviews and discussion:
<thread>
creation date:
1729777461209

thread_type:
official_review

<thread_object>
creator: reviewer_n12r

<content>
summary:
This paper introduces EM-LLM, a novel approach that integrates key aspects of human episodic memory and event cognition into LLMs with no fine-tuning, enabling them to handle practically infinite context lengths while maintaining computational efficiency

strengths:
They introduced EM-LLM, a flexible architecture that integrates key aspects of human episodic memory and event cognition into Transformer-based LLMs. Their approach enables existing LLMs to effectively process vastly extended contexts without the need for pre-training. They perform successful passkey retrieval across 5M tokens, a length which is computationally infeasible for current full-context models

weaknesses:
1. As you mentioned, almost all the methods used in your paper appear to have been proposed by others. Could you clarify what specific innovations you are introducing?

2. You stated that your method outperforms others, but in Table 1, I did not see the impressive results you described. It seems that your method may not be as effective as claimed.

3. Do you apply Eq. (1) to each token individually? If so, I imagine this would require substantial processing time. How do you address this issue?

4. In Eq. (2), could you explain how you derive the key vectors?

5. You employ k-NN in the MEMORY RETRIEVAL stage—does this significantly increase computation time?

6. I found Figure 4 difficult to understand. Could you provide additional explanation to clarify it?

questions:
The same to Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730215375916

thread_type:
official_review

<thread_object>
creator: reviewer_rkzp

<content>
summary:
This paper introduces EM-LLM, a method for extending the context window of LLMs by incorporating principles from human episodic memory and event cognition. The method segments input sequences into episodic events using a combination of Bayesian surprise and graph-theoretic boundary refinement. It implements a two-stage memory retrieval process that combines similarity-based and temporally contiguous retrieval. The authors evaluate their method on LongBench and InfiniteBench benchmarks, comparing against state-of-the-art retrieval models like InfLLM and RAG.

strengths:
The paper proposes the idea of equipping LLMs with cognitive science principles (episodic memory and event cognition), which aligns well with current challenges in long-context processing;  
The use of Bayesian surprise for dynamic event segmentation in LLMs is novel, moving beyond simple fixed-length segmentation used in prior work like InfLLM;  
The method achieves 100% accuracy on Passkey.Retrieval task with sequences up to 5M tokens, demonstrating practical scalability;  
The paper provides a scalable alternative to both RAG and full-context approaches, outperforming them on most tasks while maintaining computational efficiency, increasing its practical value.

weaknesses:
1.The paper claims surprise-based segmentation is superior to fixed-length approaches (like InfLLM), but lacks theoretical justification for why this should be effective. While Figure 4 shows some empirical correlation with human segmentation, there's no analysis of why this leads to better LLM performance. The authors should provide a theoretical analysis showing why this metric captures semantically meaningful boundaries better than alternatives.  
2.The boundary refinement process (Algorithm 1) lacks convergence analysis. Given that it's iteratively adjusting boundaries based on modularity/conductance metrics, it is important to establish guarantees about stability and convergence.  
3.The two-stage memory retrieval process (Section 3.4) introduces ks and kc parameters for similarity and contiguity buffers, but it does not justify the choice of values or analyze their trade-offs.  
4.Although Table 1 shows some improvements over InfLLM, the improvements on most tasks may not be statistically significant.  
5.The human correlation study in Section 4.2 relies on 3 short podcasts, which is insufficient to make broad claims about alignment with human event segmentation. The authors should expand the section to a larger, more diverse dataset of human annotations.  
6.Section 4.4 discusses the impact of contiguity but does not properly ablate different buffer sizes or analyze the trade-off between similarity and contiguity buffers.

questions:
1.How sensitive is the method to the choice of the surprise threshold parameter? What is the process for tuning this parameter?  
2.How do you ensure the stability of the boundary refinement process? Are there cases where it fails to converge?  
3.Does the issue of needing to query too many matching events still exist when the context window is too long?  
4.If we use the Surprise-based method to segment the context into a RAG format, will the performance remain the same?  
5.How does the method handle dynamic changes in context relevance over time?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730596236778

thread_type:
official_review

<thread_object>
creator: reviewer_pmuj

<content>
summary:
The paper introduces EM-LLM, a novel architecture that enhances the memory capabilities of LLMs by integrating design elements of human episodic memory. EM-LLM employs surprise-based event segmentation, similarity-based, and contiguity-based retrieval to enable LLMs to manage long contexts with computational efficiency. Tokens are organized into episodic events by dynamically detecting boundaries using a surprise metric and further refines these boundaries with graph-theoretic measures. This episodic structure allows EM-LLM to recall information from both the current context and prior temporally contiguous events. Experiments on long-context benchmarks show that EM-LLM outperforms InfLLM and RAG. Additionally, the segmentation results are aligned with human-perceived event structures, suggesting EM-LLM as a computational model of human episodic memory.

strengths:
Originality: EM-LLM’s incorporation of episodic memory features, particularly surprise-based segmentation, similarity-based, and contiguity-based retrieval, is a novel approach within the domain of LLMs. Using LLM surprise to model event cognition and to approximate episodic memory is original for cognitive modeling. This work not only extends the long-context capability of LLMs but also bridges machine learning and cognitive science, presenting a unique computational framework for studying memory and event cognition.
Quality: The paper demonstrates EM-LLM’s performance on recognized long-context benchmarks. The experiments highlight the model’s improvements in accuracy and efficiency over competing methods, including InfLLM and RAG.
Clarity: The paper clearly presents the concept of episodic memory in LLMs, with a logical breakdown of memory segmentation, refinement, and retrieval mechanisms. The figures generally illustrate the processes well.
Significance: EM-LLM has potential significance in advancing LLM capabilities for long-context tasks, a critical area of LLM development. Beyond practical applications, EM-LLM contributes to interdisciplinary research, providing insights into cognitive science by modeling elements of human episodic memory within an LLM framework.

weaknesses:
- In Table 1, the performances of S, SM, and SM+C methods are difficult to directly compare since different base LLMs are used in each row. I know the comparison is (partially) in Fig. 4-7 but they are not mentioned in the main text. It would be helpful if explicitly mentioned/referenced in the main text.  
- Although InfLLM is a primary benchmark comparison, it is absent from Figure 1. Including InfLLM in Figure 1 would provide a more complete comparison.
- It’s unclear from the current paper how the number of retrieved events and the number of evicted tokens influence performance. Any plot showing the model’s performance as a function of these parameters would be helpful.
- In Formula (1), notations (\mu_{t-\tau:t}) should be corrected, keeping consistent across the context.
- Certain figures would benefit from clarification, as detailed in questions.

questions:
- Can the authors discuss the feasibility and potential advantages of end-to-end training of networks for event segmentation, refinement, and retrieval, compared to these hand-designed algorithms? As a comparison, these algorithms are implemented by neural circuits in the brain.
- Can the authors discuss how might other features of human episodic memory, such as the hierarchical organization of events (human episodic memory structures can have multi-level segmentation of experiences), benefit EM-LLM’s design in terms of performance? 
- Clarifications for Figures:
In Figure 1, does the data reflect results from S, SM, or SM+C? 
In Figure 2, the memory units seem to have variable sizes, different from the fixed size in InfLLM? Additionally, how are the bars normalized?
In Figure 4A, S appears closer to human segmentation than SM or SC, contrasting with Figure 4B. 
In Figure 4, what does SC stand for? Does "C" refer to the continuity buffer or represent something else?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730652795503

thread_type:
official_review

<thread_object>
creator: reviewer_sytr

<content>
summary:
In this work, the authors proposed a framework, termed EM-LLM, for LLMs to store previously seen sequences of tokens as memory and then utilize the stored information to improve the performance of newly encountered problems. Specifically, inspired by human episodic memory, the authors suggested a method based on Bayesian surprise and graph-theoretic refinement to parse the sequences into events. Later, these store events were retrieved depending on content similarity and temporal contiguity. Experiments were conducted to demonstrate the performance improvements. In addition, a comparison was conducted to show that the parsing of events by the proposed method was similar to the results of humans, suggesting a potential link to the underlying mechanism in the brain.

strengths:
Memory is a key function in biological brains that helps organize experiences and use them to guide future behaviours.  The lack of memory function in the  LLM is a major aspect that needs to be addressed. The authors have proposed a framework inspired by the biological memory system for LLMs, which is desirable and may stimulate further investigations in this direction.  Overall, the paper was clearly written, showing comprehensive experimental results.

weaknesses:
"Human-like Episodic Memory" is an overstatement. First, episodic memory combines multimodal information into a coherent recollection of experiences, with key aspects including when, where, what, how and with who a event happened. To store a sequence of tokens does not necessarily reflect such a memory. Second, parsing of events in human memory depends heavily on the content, e.g., change of location or context. The Bayesian surprise measure may be a useful proxy to approximate such separation, but it does not by itself suggest a common mechanism, nor support the claim of "human-like". Third, the retrival rule of similarity is generic for memory, not specific for EM.

Another concern is that the memory-based frawmork has been previously proposed (cited as Xiao et al, 2024a in the current paper). The present work is an refinement of the framework, which somewhat limits its noverty.

questions:
1. I would suggest the authors to better explain the differrence in comparison with Xiao et al. 2024a, and explain why it is an important conceptual advance, instead of an incremental improvement. 
2. To better justify the claim that the mechanism is EM-related with further analyses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732039649057

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers

comment:
We sincerely thank all reviewers for their thoughtful feedback and valuable suggestions. We have provided detailed responses to each point raised, and we are currently incorporating these discussions into an updated version of the paper. As noted in our individual responses, some content will be added to the supplementary material due to space constraints. We plan to update the paper PDF later in the discussion period to incorporate any additional feedback that may arise. We believe these revisions will strengthen the paper and clarify important aspects of our work.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732309698375

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Complete paper revision: New results and detailed change log

comment:
Dear Reviewers,

We sincerely thank you for your thoughtful feedback which has helped us significantly strengthen the paper. We have now uploaded a revised PDF that incorporates all discussed changes, with particular focus on the main concerns raised:

1. Empirical validation (due to points raised by **pMUj**, **RkzP**):
    * Included new experiments where:
      * We achieved **100**\% retrieval score in **10M** tokens using a single 32GB NVIDIA GPU and *LLaMA3.1-8B* (see **Fig. 1**). To our knowledge, this is far beyond any length achieved so far with transformer models in the literature.
      * We compared surprise-based segmentation in RAG format and showed it does not beat standard RAG, let alone our approach. We discuss the reasons why.
    * Added statistical significance analysis showing consistent improvements over InfLLM (p < 0.05) in Appendix **A.1**.
    * Expanded ablation studies on parameter sensitivity and buffer sizes in Appendix **D.1**.
2. Novelty and contributions (due to points raised by **SyTr**, **N12r**):
    * Clarified our three key architectural innovations: dynamic surprise-based segmentation, graph-based refinement, and contiguity buffer (see end of **Introduction** and Appendix **E.2**).
    * Enhanced discussion of theoretical foundations linking our approach to cognitive science (see Appendix **E.1**, mentioned also in the **Discussion**).
3. Technical rigour (due to points raised by **RkzP**):
    * Updated complexity analysis for boundary refinement (Appendix **C.1**).
    * Enhanced discussion of k-NN retrieval efficiency (Appendix **C.2**).
    * Expanded parameter tuning methodology (Appendix **D.1**).
4. Clarity improvements:
    * Updated most figures and captions for better interpretation.
    * Strengthened connection between main text and appendix content.
    * Refined terminology where needed for precision.

We have also edited our responses to add details on the specific changes made for each point raised for your convenience. The results continue to demonstrate state-of-the-art performance on long-context tasks while maintaining computational efficiency. We remain fully committed to addressing any additional feedback you may have, whether about terminology, methodology, or experimental validation. We will promptly incorporate your suggestions and upload new revisions as needed.

Best regards,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
(a) This paper introduces a novel architecture, EM-LLM, designed to enhance the long-context capabilities of Large Language Models (LLMs). The authors draw inspiration from human episodic memory and event cognition to develop a system that can handle theoretically infinite context lengths. The key scientific claims include:
- Dynamic Event Segmentation: EM-LLM segments text into meaningful episodic units using a measure of Bayesian surprise, departing from fixed-length approaches like InfLLM.
- Graph-Based Refinement: The initial segmentation undergoes refinement based on graph-theoretic metrics, optimizing the cohesion within and separation between events for better information retrieval.
- Two-Stage Memory Retrieval: EM-LLM employs a retrieval process that combines similarity-based search with a mechanism mimicking the temporal contiguity effect observed in human recall, enabling more nuanced access to information.
- Good performance on LongBench and ∞-Bench.

(b) Strengths
- The proposed EM-LLM architecture offers a novel approach to long-context LLM processing, combining established techniques (Bayesian surprise, graph-theory, k-NN retrieval).
- The demonstrated ability to process vast text sequences (10 million tokens) with high accuracy showcases the model's scalability and potential for real-world applications where context is crucial.

(c) Weaknesses:
- Limited Human Correlation: While the study attempts to connect EM-LLM to human event perception, the analysis relies on a relatively small dataset (3 podcasts) and lacks deeper exploration of the nuanced relationship between model behavior and human cognitive processes.
- The initial framing as "human-like" episodic memory was considered an overstatement, prompting the authors to revise the terminology to "episodic memory-inspired".
- Some reviewers desired a stronger theoretical grounding for certain design choices, such as why surprise-based segmentation should lead to improved performance compared to fixed-length methods.

(d) The decision is to Accept.
While acknowledging the weaknesses, the novelty of the architecture, the promising empirical results, and the potential for advancing both LLM capabilities and computational models of cognition make this paper a valuable contribution.