Human-inspired Episodic Memory for Infinite Context LLMs
Authors: Zafeirios Fountas
Martin Benfeghoul
Adnan Oomerjee
Fenia Christopoulou
Gerasimos Lampouras
Haitham Bou Ammar
Jun Wang
Venue: ICLR 2025 Poster (2025)
Keywords: large language models
long context
retrieval
episodic memory
event cognition
training-free
Primary area: foundation or frontier models, including LLMs

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.

Published as a conference paper at ICLR 2025
HUMAN-INSPIRED EPISODIC MEMORY FOR INFINITE CONTEXT LLMS
Zafeirios Fountas', Martin A Benfeghoul": Adnan Oomerjee Fenia Christopoulou', Gerasimos Lampouras', Haitham Bou-Ammar and Jun Wang"
Huaawe Noah's Ark Lab, London, UK
AI Centre, Department Computer Science, University College London, London, UK (zafein sefeirios-.o fountas, adnan. ebrahim. oomer jeei @huawei com (gerasi mos. ampouras, haitham, ammar @huawei. ccom
martin. ontoine.......o hww..aatt....e...eeoeeueeueeeee httpghhhuuhh--attteee com
jun, wang@ucl. ac.uk
ABSTRACT
Large language models (LLMs) have shown remarkable capabilities, but still struggle with processing extensive contexts, limiting their ability maintain coherence and accuracy over long sequences, In contrast, the human brain excels at organising and retrieving episodic experiences across vast temporal scales, spanning lifetime. In this work, we introduce EM-LLM, novel approach that integrates key aspects of human episodic memory and event cognition into L.LMs with no fine-tuning, enabling them to handle practically infinite context lengths while maintaining computational efficiency. EM-LLM organises sequences tokens into coherent episodic events using combination of Bayesian surprise and graph-theoren boundary refinement in an online fashion. When needed, these events retrieved through two-stage memory process, combining similarity- based and temporally contiguous retrieval for efficient, human-inspired access to relevant information. Experiments on the LongBench and -o-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- perceivee events, suggesting parallels between this artificial system and its biological counterpart, thereby offering a novel computational framework for exploring human memory mechanisms.
INTRODUCTION
For contemporary pre-trained large language models (LLMs), the context window serves as the primary mechanism to incorporate domain-specific, private, or common up-to-date information. However, despite their remarkable and ever-expanding capabilities, LLMs still exhibit significant limitations when tasked with processing extensive contexts (Liu et al.) (024a)). These limitations stem from inherent challenges in Transformer-base based architectures. Recent studies have shown that Transformers struggle with extrapolating to contexts longer than their training window size (Kazem- nejad et al. 2024). On top of this, employing softmax attention over extended token sequences requires substantial computational resources for each token generation, while the resulting aggregated embeddings (the weighted sums of value vectors) risk becoming excessively noisy and losing their distinctiveness (Tworkowski et al., 2023).
To mitigate these challenges, recent works have focused on retrieval-based methods, either in the form of in-context augmentation (e.g., setrivaal-augmente generation (RAG)-based techniques (Lewis) key-value pairs (KV) within
et al. 2020 Gao et al. 2024)) or via retrieval of
"Equal Contribution. Code available httpt::////// https://githubs com/ em- 11m W...-mode.

---

Published as a conference paper at ICLR 2025
individual attention heads (Wu et al. 2022 Tworkowski et al. 2023 Bertsch et al. 2023). Notably, state-of thte--of-the-ar (SOTA) performance is achieved when KV pairs are initially organised into non- overlapping segments and then retrieved together one block of sequential tokens (Xiao al., 2024a). While such techniques present interesting research avenues, we still see significant gap between the performance of LLMss in short- long-context tasks, when existing long-context architectures are employed (Liu et al. 2024a).
This work tackles the above challenges and attempts to bridge this performance gap by taking inspiration from the algorithmic interpretation of episodic memory in the human brain the memory system respon- sible for encoding, storing, and retrieving personal experiences and events. The brain makes sense of its continuous experience the real world by segmenting it into discrete episodic events (Clewett et al. 2019; Zacks 2020), which are first organised in hierar- chical and nested- nested-timescaae timescale structure (Bal. dassano et al., (0177) and then stored in long-term memory. Notably, the bound- aries between such events the access points for memory retrieval (Michelmann etal. 2023a) and are widely believed cor- respond to points in time with high predic- tion errors between the brain's isin's generative model and its raw sensory input (a.k.a,, sur- prise). In this context, surprise refers mo- ments when the brain's 's predictions about incoming sensory information are signifi- cantly violated, leading to a mismatch be- tween what is expected and what actu-
100
Summarisution
Multi-doc QA
RAG
Full context
EM-LLM
36.44
39.3
51.58
40
Few shot learning
Single-doc
QA
Retrieval
Coding
100
Retrieval Performance
EM-LLM
6k
Mistral-7B
8k
LongLoRA
100k
Codel. aMA
100k
20-
YaRN
128k
LongRoPE
2048k
2'k 2 2
10M
Sequence Length
ally perceived. These instances of high surprise are thought to signal important narrative, changes in the environment prompting the brain to segment the ongo- ing experience into distinct events (Zacks et al. 2007| 2011 Roseboom et al. 2019
Figure 1: Top: EM-LLMs (surprise only) vs. RAG (NV-Embed-v2 retriever) vs. full-context, with LLaMA- 3.I-8B as the base LLM, evaluated on LongBench. Bot- tom: Comparison of various long-sequence methods (sorted based on their context window length) on an
extended version of oc-Bench's 's CeetreveePasskee Base-
Sinclair al. 202 Fountas et 2022
line data taken from Ding et al. (2024).
Once segmented and stored, the brain re-
calls episodic memories based on their similarity to current experience, recency, original temporal order, and their proximity to other recalled memories (temporal asymmetry and contiguity (Howard
and Kahana )002)).
Contributions: We propose EM-LLM, a novel architecture integrating crucial aspects of event cognition and episodic memory into Transformer-bast based LLMs through three key innovations (a, b and c). For memory formation, we segment input token sequences into memory units representing episodic events. The boundaries of these units are (a) initially determined using the model's 's surprise level during inference, then (b) refined to maximize within-unit cohesion and cross-unit separation (see Section 3.2). This refinement leverages graph-theoretic theoretic metrics, treating attention key similarity as a weighted adjacency matrix, and aims to enhance efficient information recall in complex, long- context tasks: by consolidating related information into single units, we seek to minimize the number of units needed for event-specific recall. The resulting memory formation process is computationally efficient: surprise-based segmentation requires no additional computation, and refinement complexity is Onem), where m is typically negligible compared to the token count in long-context tasks. For memory recall, (c) our approach combines similarity- based retrieval with temporal contiguity and asymmetry mechanisms, building recently discovered parallels between LLMs and human sequential information retrieval patterns (Ji-An et al, 2024). This method therefore ensures efficient information access while replicating temporal dynamics from human free recall studies (Howard and
2

---

Published as a conference paper at ICLR 2025
Kahana 2002).
and enhancing performance on tasks requiring temporal reasoning. See Appendix E.2
for analysis of EM-LLM's urchitectural contributions.
Performance: We show that our method is scalable and significantly outperforms the SOTA retrieval model InfLLM (Xiao et al. (024a), as well as RAG and full-context methods, on the widely-used LongBench (Bai et al. (023) and oc-Bench (Zhang et al. 2024) benchmarks designed for long-context tasks (see Fig. D. Furthermore, we perform successful passkey retrieval across 10M tokens, a length which is computationally infeasible for current full-context models. To further prove hypotheses, we then employ series of numan-annotater podeast scripts to show that information in LLM attention heads can be semantically grouped in a way that correlates with the event structure perceived by humans. Therefore, LLM-peeceived surprise can indeed serve cognitive signals that drive human event segmentation, as confirmed by previous studies (Kumar et al. 2023). Finally. using the long-context PG-19 dataset (Rae et al., (2200,, which comprises diverse corpus of English books, we evaluate the effectiveness of our segmentation method for grouping relevant information and assess the performance of different boundary refinement objectives.
proxy for the
RELATED WORK
2.1
LONG-CONTEXT IN LLMS
Recently, several approaches have been proposed to extend the context window of Transformer- based models. These include methods that address the limited representational capacity of softmax attention, and its quadratic computational and memory cost (Katharopoulos et al., 2020 Munkhdalai ct al. 2024). Other methods target the poor extrapolation of typical positional encodings to out-of- distribution context lengths (Kazemnejad et al. 2024). The latter is evident in most widely used methods, including the original absolute positional encodings (Vaswani et al. 2017) and the more
recent relative positional encodings, such as the Rotary Positional Embeddings (RoPE) (Su et al.
2024). To address this, some methods propose scaling of the rotation angles (Chen et al., 2024a) or
the base constant in RoPE (Xiong et al. 2023 Liu et al., 2024b Peng et al. 2024; Ding et al. 2024).
Others, scale positions without affecting the embedding function Press et al.
2021 Chen et al. 2023
Jin et al. 2024), explore alternative strategies such as KERPLE (Chi et al. 2022) and FIRE (Li et al.
2024a) or adopt relative position mechanisms from certain LMs like T5 (Raffel et al.
2020).
Concerning computational efficiency and diluted attention, successful approaches propose methods for general improvements to Transformer efficiency through optimised computations (Dao) 2024 Han
al. 2024a Aminabadi al.. 2022 Kwon al.) 20233] Liu et al., [202c] Brandon al. (023) or
compression techniques (Nawrot et al. 2024| Zhang et al. 2023). as well as training methods tailored
for long-context scenarios (Zhu et al.) [024]] Chen et al. |l.|||122210 2024b). Another direction is the utilisation of
retrieval- methods, the vast majority of which relies on vector database that keeps key-value cache and scalable approximations of k-nearest neighbors (k-NNs) to perform lookups (Wu et al., 2022, Tworkowski et al.) 2023) [Bertsch et al. 2023). Interestingly. since using key-value cache with k-NN lookup can be seen as an approximation of applying softmax attention to the full token sequence (see Appendix FFII.. k-NN retrieval methods can be used without fine-tuning (Bertsch et al., 2023). For an exception that does not rely on k-NNs, see Wang et al. (2023).
recent and interesting variant of k-NN retrieval involves retrieving large groups of tokens, rather than individual ones. Models that rely on this approach include SLED (Ivgi et al. 2023) and the more recent InfLLMM (Xiao et al. 2024a), which achieves SOTA performance on long-conted benchmarks. InfL.MM segments the entire context length into fixed-size memory units and employs k-NN lookup using the tokens with the highest accumulated scores per unit. The latter can be seen as a form of hierarchical attention in models that use such retrieval, as illustrated in Fig. 2 While group- based retrieval represents a promising direction, our approach significantly advances this concept by dynamically determining token groupings in manner akin to human memory formation. This effectively addresses a fundamental limitation of InfLLM's fixed-size segmentation and enables more adaptive and context-sensitive processing of extended information.
2.2 NEURAL MODELS OF EPISODIC MEMORY AND EVENT COGNITION
The concept of episodic memory, central to our approach, has been extensively studied in both theoretical neuroscience and machine learning. Neural models of episodic memory capture human behaviour and neuroimaging data, providing insights into how the brain processes and stores expe- riences and suggesting links between memory, efficient representations and navigation of physical
3

---

Published as a conference paper at ICLR 2025
Level episodic attention
Lavel 2: softmax attention on selected groups (contiguios view)
Simihitt score
25-
Group
50
100
150
200
250
20
30
40
50
70
80
All post tokzess (posttion id)
Connect window (position kd)
Figure 2: Group-based k-NN retrieval can be seen as form of hierarchical episodic attention. Initially, k = 4 groups of tokens are selected (left) and then used for softmax attention (right), as if all other similarity scores were forced to be zero mhon-shaded areas of the left curve). This framework can support multiple levels of episodic attention.
and conceptual spaces (Gershman al. 2012 Benna and Fusi 2021). In machine learning, episodic
-eepory-inspired approaches have yielded significant improvements across various domains. For instance, episodic control has enhanced reinforcement learning agents' performance and learning speed (Blundell et al.) [016)] PPritzel et al., 2017 Coda-Forno al. (022). In addition, models of
memory construction and consolidation have been successful alleviating catastrophic forgetting
in neural networks (Kirkpatrick et al. 2017 Lopez- and Ranzato 2017: Chaudhry al. 2019
Buzzega et al.) 2020; Prabhu et al. 2020). including LLMs (Das et al. 2024). ) and appear to cxp explain
key features of human memory, such as imagination and future thinking (Spens and Burgess 2024).
These models have revealed key aspects of episodic memory, particularly in describing how experi- ences are segmented into events, and when new memories are encoded and retrieved (Lu et al. 2022).
Surprise plays critical role in this process, triggering event boundaries and memory formation (Foun
las et al. 2022 Kumar et al.] (002)). This event-based structure is deeply intertwined with our perception of time (Roseboom et al. 2019 Sherman et al. 2022), highlighting the interdependence of memory and temporal cognition. This insight has helped generative models for video (Zakharov
et al. 2022a b) and reinforcement learning (Zakharov et al. 202T) to capture temporal dynamics more accurately. In terms of memory retrieval, studies in human free recall have shown a distinctive increased likelihood of retrieving items encoded close together in time (temporal contiguity) and in succession (temporal asymmetry) (see Fig. [ Recently, it was shown that attention heads in Transformer-based LLMs that are associated with in-context learning, already exhibit the same dynamic retrieval behaviour (Ji-An et al. 2024) (Fig. 3B) including both contiguity and asymmetry effects. Therefore, Transformers have the inherent ability to act as episodic memory retrieval models, if provided with the right information within their context window. Our work leverages these concepts of surprise-baser based event segmentation and LLMs' inherent temporal contiguity and asymmetry effects to enable new generation of Infinite Context-Length LLMs, capable of processing and understanding information over vastly extended timescales.
EM-LLM: LLM WITH EPISODIC MEMORY
3.1
ARCHITECTURE
EM-LLM is designed to be applied directly to pre-trained LLMs, enabling them to handle context lengths significantly larger than their original training length. Our architecture, illustrated in Fig. 3 divides the context into three distinct groups: initial tokens, evicted tokens and local context. This structure, while incorporating insights from recent work on token block retrieval (Xiao et al., 2024a). introduces novel elements inspired by human episodic memory.
The local context represents the most recent tokens, maximising information about the current task, and fits within the typical context window of the underlying LLM. This group utilises full softmax attention and plays a role similar to the focus of attention in cognitive models of working memory, holding the most immediately relevant information for the current task (Cowan) 200T). The evicted tokens typically comprise the majority of past tokens in long-context scenario, extending far beyond the LLM's s original training length. These tokens are managed by our proposed memory model functioning similarly to short-term episodic memory in the brain. Finally, following previous work, we also maintain a group of 128 initial tokens in the LLM context. These act as attention sinks and help recover the performance of window attention, first observed by Xiao et al. (2024b); Han ct al. (2024b) and later adopted by Xiao et al. (2024a). For retrieved tokens, which are therefore discontinuous and outside the local context, we assign fixed position embedding as in Raffel
4

---

Published as a conference paper at ICLR 2025
A
surprise
Conditiveal resporose probubiitt
122
eVEnT
event
initial tokens
episodic memory
local context
Distance from recalled item
Avenge attention sore
contigious event
contigious event
deguree
eaquene
initial tokens
contiguity butfer
similarity buffer
local context
Distance from recalled token
Figure 3: (A) Example of the temporal contiguity and asymmetry effect in human free recall. Data
averaged over several large free recall studies (adopted from Howard and Kahana (2002)). (B) The attention scores of a GPT2 head averaged over all tokens tested (adopted from Ji-An al. (2024)
(C) Schematic illustrating our proposed process for memory formation and retrieval in each layer: O Input sequence with surprise-based segmentation (purple arrows indicate high surprise). @ Formation of episodic memories; input is segmented into events and stored, with initial tokens and local context preserved. Note that the boundary refinement process not shown here for clarity. @ Memory retrieval via k-NN search, selecting contiguous events from episodic memory. + Final context window structure, comprising initial tokens, contiguity buffer (populated by neighbouring events), similarity buffer (from k-NN retrieval), and local context.
al. 2020): : Xiao et al. (002a).. This architecture enables EM-LLM to effectively process and utilise information from positions outside its pre- pre-trained local context window, while maintaining the underlying LLM's performance characteristics.
3.2 MEMORY FORMATION VIA SURPRISE
In the context of LLMs, we define episodic memory as the organised, event-based collection of past key-value pairs, analogous to the latent representations of personal experiences in human memory. Just as unexpected or novel information plays a crucial role in human memory formation, we posit that analogous indicators of novelty in LLMs can serve effective proxy for identifying significant "events" within the model's experience. In Bayesian terms, surprise is quantified by the negative log-likelihood of observing the current, ground-truth token given the previous tokens in auto- regressive model, with high values indicating the unpredictability or novelty of each new token within the context according to the model, i.e., being "surprised" by the next token. Following work on cognitive modelling (Roseboom et al.) 2019, Fountas et al. (022), we employ thresholding mechanism to perform an initial identification of event boundaries (used for the first time in LLMs). Formally, a token x, is considered a potential boundary if its surprise value exceeds threshold T:
- log P(x; I 2t-1; 0) > T with Pt-rt + Y0t-7t
(I)
where P4-rat and of-rst are the mean and variance of surprise for a window offset T, and 7y is scaling factor. The choice of threshold T' is critical in balancing the granularity of segmentation with the model's sensitivity to contextual shifts. If the T is too high, we will identify very few event boundaries, especially if the local context contains few surprising tokens. Conversely, a low T results in frequent boundary identification. Using a moving window ensures that T adapts to contextual shifts, minimizing the need for manual tuning while maintaining control over threshold sensitivity via This initial segmentation results in set of potential event boundaries B = b,, b2,..., ba. where each b; represents the index of token exceeding the surprise threshold. These boundaries serve as the starting point for our subsequent refinement process, which aims optimise the intra-event coherence and inter-event distinctiveness of the resulting memory segments.
3.3 BOUNDARY REFINEMENT
While surprise-bar based segmentation provides an effective initial estimate of event boundaries, we make the key observation that the utility of elements within event, during memory recall, depends
5

---

Published as a conference paper at ICLR 2025
Algorithm Event segmentation in KV cache
Input: tok: List of tokens in the sequence
Input: T: Threshold for surprsaa to identify initial boundaries
Input: f: Metric function to evaluate potential boundaries
Output: B: List of final boundary positions
I: B <- [i for i in range(length (tck)) if log (og(P(toklii)) 71
Boundary identification
2: for i in range(length(B) do
ax, B +
4:
+| f- arg max max BEta. f(A,(a,8))
Boundary refinement
5: end for
6: return B
on their likelihood of being utilised by the current query. Therefore, we theorise that memory recall will be most efficient with high intra-event similarity between keys while maintaining low inter-event similarity. For instance, see the similarity of groups in Fig. 2 To further ensure this, we introduce boundary refinement step that looks to optimise this objective. Such an objective is typically optimised in the context of graph-clustering, hence we express this refinement process in graph-theoree theoretic manner. To achieve this, we treat the similarity matrix between all keys of an attention head h within the local context window for tokens Z1, 22, In as an adjacency matrix A":
4 s KY))
(2)
where K. and Ky are the key vectors corresponding to tokens x; and æy, respectively. The similarity function measures the closeness of two key vectors; in our implementation, we use dot product similarity 1hh T KP due to its effectiveness in capturing semantic relationships in high-dimensional spaces (Vaswani et al. 2017) and to align with the mechanism of self-attention in Transformers.
To evaluate the quality of potential boundaries, we define a metric function (AAB)) TR=x x (I. E. This function quantifies the cohesion within events and separation between events based on the graph structure represented by the similarity matrix A and event boundaries B. We experiment with two widely-accepted graph-clustering metries: modularity and conductance (Mi- asnikof et al. 2018). Modularity (Newman and Girvan 2004) provides a measure of the quality of a particular division of a network into communities, with higher values indicating higher edge density in the identified cluster when compared to the density of edges expected in random cluster. As our edge weights represent the similarity between two tokens, we seek to maximise this metric. Modularity is defined as:
,, 8(c,,
far (Ar 4m
2m
where m is the total edge weight in the graph, is the community (cpisodic event) to which node i assigned, and 8 is the Kronecker delta function. Conductance, on the other hand, measures the fraction of total weighted edges cut by given community boundary, and is defined as:
(A A,, BB) = min EiES,Jes Al SEV min (vo(),, vo(V S)) with vo(S) = Ag, vo(V () = Σ A (4) JJES
iJsS
where = (b, by 1..., b subset of all nodes = (b,, bi 1,..., by) in the induced graph, with b; B. Lower conductance values indicate better community structure. Our boundary refinement algorithm sequentially adjusts the initial surprise-based boundaries to optimise these metric functions. While our best results are achieved using modularity, we also include comparisons with conductance-bast
boundary refinement to provide a comprehensive analysis. The overall process is summarized in Algorithm [ and further discussed in Appendix E.3
This algorithm first identifies initial boundaries based on the surprise threshold T, then refines these boundaries by finding the optimal position / between each pair of consecutive initial boundaries (a, 3) that optimises the chosen metric function f (either maximising modularity or minimising conductance). This process ensures that the final segmentation (1) captures points of high surprise and (2) optimises for coherent information grouping. The boundary identification step incurs negligible computational cost, as it only evaluates existing LLM outputs. The time complexity of Algorithm []] has an overall complexity of (nm)), where n is the n is the sequence length and m is the chunk size selected to process the sequence (for details see Appendix C.1

---

Published as a conference paper at ICLR 2025
Base LLM
Method
LongBench Sum
Bench
SQA
MQA
FSL
Ret
Cod
Avg.
C.D
M.F
MC
R.KV
R.P
R.N
Mistral
InfLLM (4k+2k)
33
25.5
27.1
66.1
64
54.8
41.9
29.4
26.6
43.2
95.6
100
99.8
EM-L.MMMM-C
32.9
27
27.2
66.8
84.1
54.8
43.7
28.2
27.1
42.8
99
100
99.8
I.LaMA
InfLLM (4k+4k)
38.5
36.9
69
84
53.2
47
30.5
23.7
43.7
5
100
99
3
EM-LLMs
39.3
37.7
27.0
69.2
87.5
50.3
47.2
31.7
16.9
40.6
4.2
100
99.5
LLLMA
InfLLM (4k+4k)
41.4
40.7
29
69
64.2
51.1
22.6
33.7
46.7
81
100
100
3.1
EM-LLMsM
41.2
41.3
29.2
69.1
98.5
64.1
51.3
22.6
34
47.6
90.2
100
100
InfLLM (Ik+3k)
28.4
24.9
25.6
52.9
7.5
57
34.5
Phi 3
EM-LLMs
29.2
27.1
25.9
53.5
10
57
35.4
InfLLM (Ik+3k)
31.7
28.5
23.9
56.3
11.5
40.3
34.2
Phi 3.5
EM-LLMs
31.8
31.9
24.5
55.5
13
39.5
34.9
Table 1: EM-LLM performance on LongBench (grouped tasks) and co-Bench compared to baseline InfLLMM. S: surprise threshold, SM: surprise threshold and refinement with modularity, S+C: surprise threshold and contiguity buffer, SM+C: surprise, refinement and contiguity buffer. Each row indicates the number of local + retrieved tokens (eg. 4k+2k) used for both InfLLM and EM-LLM. See Appendix D.I for parameter choices and Appendix A.I for more results and significance testing.
3.4
MEMORY RETRIEVAL
When inferring a new token, a number of episodic events are selected and become a part of the (extended) context window of the underlying LLM. Our memory retrieval process employs two- stage mechanism to select relevant episodic events for the LLM's context window (Fig. 3C) First, retrieve k2 events using k-NN search based dot product similarity between the current query and representative tokens of each event. These representatives, selected as per Xiao et al. (2024a), are the most influential tokens within each event. For large memory stores, we utilise approximate k-NN (Douze et al.) 2024) to maintain efficiency. These k, events, retrieved based on their similarity the current query, form a part of the LLM's context window that we refer to as the similarity buffer.
The second stage of our retrieval process introduces another buffer, which we refer to as the contiguity buffer, designed to maintain temporal context. Implemented as a queue of size kec this buffer promotes temporal relationships in retrieval. When event is retrieved, we also enqueue its neighboring events (within in positions in the original sequence) into this buffer. This mechanism enables the LLM's "induction" attention heads to exhibit the contiguity and asymmetry effects discussed in Section 2.2 The queue structure allows for natural decay of temporal context as new events are processed, with older or repeated events being dequeued as new ones are added. In total, = k2 ke events are added to the context window, striking a balance between relevance and temporal relationships in manner analogous to human episodic memory retrieval, Note that each layer retrieves and attends to these h events individually, allowing it to potentially focus on different parts of the context.
4 EXPERIMENTS
4.1 PERFORMANCE OF EM-LLM ON LONG-CONTEXT TASKS
Comparison with -V-retrieval-basee LLMs At the time of writing, InfLLM is considered to be the SOTA KV-retrieval method on long-context benchmarks (LongBench, oo-Bench), as well as being the only method which uses group-based k-NN retrieval in LLMs such benchmarks. We, therefore, employ this model as our first baseline for comparison with our own methods.
Results on both benchmarks (Table I show that our method is able to improve on InfLLM across 5 different base LLMs, 80% of individual task groups of LongBench and on the overall average. Note that the table shows the best single method terms of overall performance for each ablation (see Appendix[A.1] for all ablations in methods). Looking at individual task performance across all ablations in methods, EM-LLM is able to surpass InfLLM in all tasks. Notably, we see an especially large jump in performance in the retrieval (Passage, KV, Passkey, Number) and QA (Narrative, Qasper, MultiField, Hotpot, Wiki and Musique) tasks across all ablations, with up to a 40% and 29.7% improvement over InfLLM respectively. Such tasks require the model to identify and retrieve specific information within the input sequence, a challenging test for the model's 's ability to accurately recall wide range of detailed information from large context concurrently. This substantial
7

---

Published as a conference paper at ICLR 2025
A
B
1.04
1.5
FM FC
SM
Intre- lntte-simillere
Modularity
SC
L.03
Humass
1,02-
FM
0.46
144
SM
9.25
1.00
FC
0.25
1.3
1.00
SC
0.19
15
20
25
layer
layer
Wasserstein distance human reports (x10)
Figure 4: Comparison of human event segmentation with different computational segmentation methods human-annotated audio dataset (see also Appendix BB). (A) Difference in metrics for the cohesion and separation of KV cache of each LLaMA2 layer. The graphs report the difference of each method with the corresponding random segmentation. (B) Distance between human reports and different methods. In both sets of results, fixed methods (F, FM, FC with M: Modularity, C: Conductance) perform worse than their surprise-bar counterparts (S, SM, SC) with InfLLMM's method (F) performing worse than random.
improvement highlights the effectiveness of our event segmentation method in enhancing long-term memory recall and retrieval accuracy in LLMs.
Comparison with RAG and full-context LLMs To evaluate EM-LLM against prominent methods for handling long contexts, we compared its performance on LLaMA 3.1-8B with two different RAG approaches, including the current SOTA NV-Embed-v2 retriever (Lee et al. 2024), well as with the brute-force baseline of processing all tokens directly within the L.LM's 's softmax attention (full- context). Across most tasks in our benchmarks, EM-LLM outperformed both RAG and full-context methods, as well as a custom surprise-base based RAG method (Fig. [ and Appendix AA.2). exceeding the performance of NV-Embed-v2 by 30.5% LongBench and by 11.5% on cc-Bench.
This significant performance boost over RAG can be attributed to EM-LLM's ability to retrieve and incorporate relevant information each layer individually, rather than relying on single retrieval step as in RAG (for an illustration, see Supp. Fig. [ By accessing more specific and contextually relevant information through layer-wise key-value retrieval, EM-LLM effectively addresses RAG's limitations in precision and lower overall performance (Li et al. 2024b). Additionally, EM-LLM's hicrarchical attention avoids the issue diluted attention in large context windows that affects full-context models, enabling it to outperform both RAG and full-context LLMs on the LongBench dataset. Furthermore, EM-LLM demonstrated remarkable scalability by achieving 100% accuracy on the Resskey.Retrieval task with sequences up 10.2M tokens, far beyond the practical limits Retrieval
of full-context LLMs. This highlights EM-LLM's efficiency in handling extremely long contexts, positioning powerful alternative for long-context processing.
4.2 HUMAN AND LLM SURPRISE CLUSTER SIMILAR TOKENS TOGETHER
As mentioned in Section 3.2 we employ modularity and conductance as two refinement objectives in our boundary refinement algorithm, due their qualities in assessing the intra- and inter-event similarities between individual tokens. We will now use such metrics to compare various event segmentation methods, including human event segmentation data. Additionally, we introduce one further, simple metric for this experiment: the ratio between intra- and inter-community similarity (I/IS), calculated for each head and community S as follows:
intra
intra
Σ Agj: iES,jES
inter
Σ Auj. VESJEES
I/IS
(5)
inter
Kumar et al.] (2023) found strong correlations between human-perceivel events and prediction errors across 3 short podcasts (7-30 minutes), when processing the corresponding transcript with an LLM. Taking advantage of such human-annotated data and results from previous studies on this dataset (Michelmann et al. 202T Lositsky et al. 2016), we compare the segmentation quality and correlation with human segmentation for each of our methods (Fig. 4 using our similarity metrics.
As shown in Fig. 4A human- perceeved events achieve significantly higher scores in similarity metrics compared to fixed or random events. suggesting that surprise is indeed an important factor for humans in their own perception of events. Furthermore, surprise-only segmentation (S) achieves very similar results to humans, while the addition of our refinement algorithm (SM, SC, FM, FC) significantly
8

---

Published as a conference paper at ICLR 2025
LLM
Metric
F
FM
FC
S
SM
SC
Mod f
-2.3 4.1
29.2 + 44.0
6.7 = 25.9
18.6 + 29.6
39,9 55.5
29.5 42.7
Mistral- 7B
Con 4
9.7
+16.9.66.7
-12.5 9.6
-23.6 9.4
-24.6
-27.6 9.8
I/IS f
-4.3++ 4.0
31.2
3.7 + 14.9
17.9 + 17.0
35.3 + 27.7
21.6 22.4
Mod
4.3
13.4 19.5
0.6
8.7
18.7 = 26.4
11.5 = 19.4
LLaMA2- 7B
Con 4
1..9
-18.8
-13.7 10.9
-29.5 10.2
-29.7 - 10.1
-33.3 + 10.3
I/IS
+3.8
20.7 184.7
-1.1
15.0 880.0
25.0 + 19.9
16.5 + 15.4
Mod f
31.6
25.9
0.9 11.8
1
27.0 = 35.6
18.3
LLaMA3-
Con L
11.3
-20.3 6.9
-4.66114 11.4
-22.7
-30.6
-33.9 +
8B
I/S f
-3.8
24.5 13.9
-1.1
15.7 11.0
28.1 16.1
16.4 12.2
Table 2: Comparison with graph-theoretic metrics in the KV cache of different LLMs and segmenta- tion methods using the PG-19 dataset and y 10-3 Reported values the difference with random segmentation. Mod: modularity x 10², Con: conductance, I/IS: intra/inter-similarity x x10s.
improves performance, Fig. 4 further shows that surprise-based methods (S, SM, SC), consistently identify event boundaries that are closest to those perceived by humans.
4.3 COMPARING SEGMENTATION METHODS
Our experiments on the PG-19 dataset (see Table 2) clearly demonstrate that cupprssebased seg- mentation with refinement (SM, SC) provides the best results in terms of event similarity metrics, regardless of the base LLM used. While the surprise-only method (S) achieves decent results, we observe that refinmment is especially adept to improving this performance with regards to our metrics, as it is directly optimising for such an objective. Interestingly however, the fixed-based refinement methods (FM, FC) do not reach the same performance as their surprise-based counterparts, further showing that the initial segmentation with surprise threshold crucial to achieving the best possible balance in intra-finter-similarty with our methods.
4.4 SIMILARITY, CONTIGUITY, RECENCY AND TEMPORAL ORDER
As demonstrated in Tables T and 2 along with Fig. each of our ablations show various positive improvements on InfLLM. As mentioned in Section in improving our similarity metrics. This is seen translate well to model performance in our experiments, with the addition of refinement achieving the best performance in 60% of tasks across LongBench and co-Bench (see Tables as well as agreeing with human data (Fig. 4 The effects of contiguity are also clearly demonstrated, with the addition of our contiguity buffer achieving the best performance in 44% of tasks. Furthermore, these methods are seen be complementary, often improving on both individual additions.
refinement has strong positive impact
However, the fact that certain tasks still appear to benefit more from either surprise-only, refinement, or contiguity, is an interesting result. This is likely due the nature of the tasks and the varying importance of contiguity across these tasks. Where contiguity is not crucial, adding such a buffer to our context window also reduces the size of the similarity buffer, and therefore provides potentially less directly relevant events. This is compatible with our own findings that contiguity buffer that as big or smaller than the similarity buffer yields the best results (see Fig. 13), suggesting that the similarity buffer is still the most crucial part of our approach. This is especially the case when combined with refinement, which we expect is due to the improved similarity of refined events, hence further reducing the need for contiguous events.
5 DISCUSSION
Human studies Significant correlations have been found between human event segmentation and prediction errors in both LLMs (Kumar et al.] 2023) and video models (Fountas et al. 2022 Mariola al. 2022). Our results add to this growing body of evidence, demonstrating that can serve as a proxy for human event segmentation, in multiple levels of hierarchical abstraction, and that the resulting event structure EM-LLM's attention heads correlates strongly with human- perceived events. This finding suggests potential, low-level parallels between LLM mechanisms and human cognitive processes (see also Appendix E.T). Furthermore, our model's use of both similarity-basee based and temporally contiguous retrieval mechanisms parallels human memory retrieval patterns, allowing for the expression of robust phenomena found human memory research (Howard) and Kahana 2002). The temporal contiguity effect, where items experienced close together in time
surprise
9

---

Published as a conference paper at ICLR 2025
often recalled together, is robust phenomenon in human memory research (Howard and Kahana 2002). Further experiments could deepen our understanding of the connections between EM-LLM and human episodic memory. Following MMchelmann al. (2023b), one could test whether the timing of the event boundaries or the degree of modularity per level that our method produces is closer average to the human consensus, than individual human subjects. Additionally, exploring how different ratios of the contiguity buffer affect the reproduction of human memory biases, and investigating the impact of recency and initial surprise on event recall, could reveal the extent to which EM-LLM exhibits biases found in free recall studies.
Furthermore, EM-LLM's architecture with differentiated context handling (Section 3.1 invites com- parisons cognitive models of human memory beyond episodic. The local context, holding recent and task-relevant information, resembles the limited-capacity working memory system described by Baddeley (2003). Given that EM-LLM's broader context window includes both local context and retrieved memories, it aligns more closely with Ericsson and Kintsch (9955> 's concept of long-term working memory, which allows rapid access to relevant long-tern -erm infommaiion beyond traditional capacity limits. Alternatively. our architecture parallels Cowan (2001)ss 's embedded processes model, where the local context is the "focus of attention". and the full context window represents the activated portion of long-term memory. Future work could explore these analogies further, using EM-LLM as a test-bed for hypotheses about human memory and working memory capacity limits. Inspired by Baddeley's 's multi-component model, integrating modality-specifie buffers into EM-LLM might enhance performance on multi-modal tasks.
embedded-processes
Machine learning In refining event boundaries, utilised modularity and conductance as metrics for evaluating community structure in the similarity graph of attention keys. While effective in our experiments, we acknowledge that numerous other methods for graph clustering and sequence segmentation could potentially be applied (Fortunato) 2010; Yang et al. 2016). Our choice was motivated by their established theoretical foundations and computational efficiency, though com- parative studies suggest performance can vary based on network characteristies (Yang et al. (016)) Interestingly, our surprise-based initial boundary detection shares similarities with Bayesian online change-point detection (Adams and MacKay 2007), suggesting potential avenues for integrating time series analysis techniques into LLM context processing. Future work could explore whether more sophisticated segmentation clustering algorithms could improve EM-LLM's performance, particularly for extremely long contexts or streaming data scenarios. Such investigations could enhance our model and contribute understanding how information is structured and processed in LLMs, bridging the gap between traditional sequence analysis and LLM context processing.
Looking ahead, promising directions for future research include extending our segmentation processes to operate at each layer of the Transformer independently. This could lead to more nuanced and hierarchical representations of episodic memories, following the underlying semantic structure of the input more closely. Additionally, exploring how EM-LLM could be utilised to enable imagination and future thinking has great potential for advancing model- based reinforcement learning and continual learning techniques in LLMs. By leveraging its event-based structure to simulate potential future scenarios or recall past experiences in novel contexts, EM-LLM could enhance an LLM'' ability to plan, adapt, and learn continuously from new information.
CONCLUSION
In this work, we introduced EM-LLM, flexible architecture that integrates key aspects of human episodic memory and event cognition into Transformer-basee LLMs. Our approach enables existing LLMs to effectively process vastly extended contexts without the need for pre-training, demonstrating superior performance on long-context tasks compared to the corresponding SOTA. By combining surprise-based event segmentation, graph-theorets boundary refinement, and two-stage memory retrieval process, EM-LLM offers a promising path toward virtually infinite context windows. This capability has the potential to revolutionize interactions with LLMs, enabling continuous, personalised exchanges over extended periods and serving viable alternative traditional RAG techniques. Finally, by bridging insights from cognitive science with machine learning, our approach not only enhances the performance of LLMs on long long-context tasks but also provides scalable framework for computational modelling of episodic and event cognition. We hope this study inspires the community to expand research the intersection of L.MMs and human memory.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024a.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024.
Szymon Tworkowski, Konrad Staniszewski, Mikolaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Miloš. Focused transformer: Contrastive training for context scaling. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL htttps://openrvviee
net/forum s1F jx
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, al. Retrieval-augmented genera- tion for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 33: 9459-9479 2020.
Yunfan Gao, Yun Xiong, Xinyu Gao, Kangxiang Jia, Jinliu Pan, Yuxi Bi, Yi Dai, Jiawei Sun, Qianyu Guo, Meng Wang, and Haofen Wang. Retrieval-augmentt augmented generation for large language models: survey, 2024.
Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In International Conference on Learning Representations, 2022. URL https
7/openrevia n www.eeeeeeeeooeeeeeeeeeeeeeeeeee fo
Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew R. Gormley. Unlimiformer: Long-range transformers with unlimited length input. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL
nnt. oot.:ppnreeeeeeeeeeeeeoooooomoo. forumm1d=1....J.
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Song Han, and Maosong Sun. Infilm: Unveiling the intrinsic capacity of 1lms for understanding ex- tremely long sequences with training-free memory. In Advances in Neural Information Processing Systems, 2024a. To appear.
Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending IIm context window beyond million tokens. arXiv preprint 1375 2024.
David Clewett, Sarah DuBrow, and Lila Davachi. Transcending time in the brain: How event memories are constructed from experience. Hippocampus, 2963):162-183, 2019.
Jeffrey M Zacks. Event perception and memory. Annual review of psychology. 71:165-191, 2020 2020.
Christopher Baldassano, Janice Chen, Asieh Zadhood, Jonathan W Pillow, Uri Hasson, and Kenneth A Norman. Discovering event structure in continuous narrative perception and memory. Neuron, 95 (3):709-721, 2017.
Sebastian Michelmann, Uri Hasson, and Kenneth A. Norman. Evidence that event boundaries are access points for memory retrieval. Psychological Science, 34(3):326- 344, 2023a. doi 1011770956697622122280 09567976221122
URL hatps://doll.or/220.17777
PMID: 36595492.
Jeffrey M Zacks, Nicole K Speer, Khena M Swallow, Todd S Braver, and Jeremy R Reynolds. Event perception: mind-brain perspective. Psychological bulletin, 133(():273, 2007.
Jeffrey M Zacks, Christopher A Kurby, Michelle L Eisenberg, and Nayiri Haroutunian. Prediction emror associated with the perceptual segmentation of naturalistic events. Journal of cognitive neuroscience, 23 23(12):405--466 2011.
11

---

Published as a conference paper at ICLR 2025
Warrick Roseboom, Zafeirios Fountas, Kyriacos Nikiforou, David Bhowmik, Murray Shanahan, and Anil K Seth. Activity in perceptual classification networks as a basis for human subjective time perception. Nature communications, 10(1):267, 2019.
Alyssa H. Sinclair, Grace M. Manalili, Iva K. Brunec, R. Alison Adcock, and Morgan D. Barense. Prediction errors disrupt hippocampal representations and update episodic memories. Proceedings of the National Academy of Sciences, 118651))2211766511188 2021. doi 10. 20.1077//aa..21176222.. URL
pras. http://ww.paaaooooo///o///a//a1/
Zafeirios Fountas, Anastasia Sylaidi, Kyriacos Nikiforou, Anil K. Seth, Murray Shanahan, and
Warrick Roseboom. A Predictive Processing Model of Episodic Memory and Time Perception.
Neural Computation, 34(7):5501115444 06 2022. ISSN 0899-76677 doi 162//eco___
URL htttps://6i..rggg wot.://oo/g...1//oooooooooo
Mare W Howard and Michael Kahana. A distributed representation of temporal context. Journal of mathematical psychology, 46(3):269-299, 2002.
Li Ji, Corey Y. Zhou, Marcus K. Benna, and Marcelo G. Mattar. Linking in-context learning in transformers to human episodic memory. In Advances in Neural Information Processing Systems, 2024. To appear.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023.
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, et be-bench: Extending long context evaluation beyond 100k tokens. arXiv preprint arXiv:2402. 13718, 2024.
Manoj Kumar, Ariel Goldstein, Sebastian Michelmann, Jeffiey M Zacks, Uri Hasson, and Kenneth A Norman. Bayesian surprise predicts human event segmentation in story listening. Cognitive science, 47(10);e13343 2023.
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Com- pressive transformers for long-range sequence modelling. International Conference on Learning Representations, 2020. httttss://peeeeeieeeeee wwt.leewweeeeeeee.eeeeeee..eeeee forum? Le
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. International conference on machine learning, pages 5156-5165. PMLR, 2020.
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 568:127063, 2024. ISSN 0925-2312. dat.hnns///bboooyy///ooooooonooo ccom/ Coom://cccnnee//rrccee///////222
URL htttps:/ww..ccenncee
Guanzheng Chen, Xin Li, Zaiqiao Meng, Shangsong Liang. and Lidong Bing. CLEX: Continuous length extrapolation for large language models. In The Twelfth International Conference on Learn- ing Representations, 2024a.
Wenhan Xiong, Jingyu Liu, Igor Molybog. Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309.160.33 2023.
Xiaoran Liu, Hang Yan, Chenxin An, Xipeng Qiu, and Dahua Lin. Scaling laws of baseda extrapolation. In The Tweifth International Conference on Learning Representations, 2024b. URL https::/openrert net wtt.://oooommmo110000000055
12

---

Published as a conference paper at ICLR 2025
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. YaRN: Efficient context win- dow extension of large language models. In The Twelfth International Conference on Learning Representations, 2024.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. 00008./2400 /2409, 2021.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint 0rXiv:2306.15595, 2023.
Hongye Jin, Xiaotian Han, Jingfeng Yang. Zhimeng Jiang, Zirui Liu, Chia- Yuan Chang, Huiyuan Chen, and Xia Hu. Lim maybe longlm: Self-extend Ilm context window without tuning, 2024.
Ta-Chung Chi, Ting- Han Fan, Peter Ramadge, and Alexander Rudnicky. KERPLE: Kernelized relative positional embedding for length extrapolation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL httpp:://ppsneevieeeee www.sviiwwweeefrrurttie 1d= XDwm
Shanda Li, Chong You, Guru Guruganesh, Joshua Ainslie, Santiago Ontanon, Manzil Zaheer, Sumit Sanghai, Yiming Yang, Sanjiv Kumar, and Srinadh Bhojanapalli. Functional interpolation for relative positions improves long context transformers. In The Twelfth International Confer- ence on Learning Representations, 2024a. URL hwtph:tts////. httpttttp////ottttet http::::///ppeeeveeeeee // net/ htt.://wwwwweteeeotmmmmmmam rRC3gFeagk
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. URL
httt::/3jmin www.ooogg/pp.eeee.... w21/20-122222 0001200074..50 html
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024. URL. https
/openreview wtt.:oommmmmmmmm2mmm3mmmmm3
Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David Woodruff, and Amir Zandich. Hyperattentio: Long-context attention in near-linear time. In The Twelfth International Confer- ence on Learning Representations, 2024a. URL htl.ht/s:////. /ottps::/oppereviii ew net. forum? frmmimd dd=
EhCOd2BJIM
Reza Yazdani Aminabadi, Samyam Rajbhandari, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Olatunji Ruwase, Shaden Smith, Minjia Zhang, Jeff Rasley, and Yuxiong He. Deepspeed- inference: enabling efficient inference of transformer models at unprecedented scale. In Proceed- ings of the International Conference on High Performance Computing, Networking, Storage and Analysis, SC 22. IEEE Press, 2022. ISBN 9784665454445.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, SOSP '23, page 611-626, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400702297. 05//T001145/6/0006666636666 URL httph://o//oiii org/
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ringattention with blockwise transformers for near- infinite context. In The Twelfth International Conference on Learning Representations, 2024c.
URL [RR/htpps///ppnrevvee net. aoteee/oooommii-isssiiiii
William Brandon, Aniruddha Nrusimha, Kevin Qian, Zachary Ankner, Tian Jin, Zhiye Song, and Jonathan Regan-Kelley. Striped attention: Faster ring attention for causal transformers. arXiv preprint arXiv:2311.0943/ 2023.
Piotr Nawrot, Adrian Lańcucki, Marcin Chochowski, David Tarjan, and Edoardo M Ponti. Dynamic memory compression: Retrofitting Ilms for accelerated inference. arXiv preprint ox31v22403.09666, 2024.
13

---

Published as a conference paper at ICLR 2025
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Re, Clark Barrett, Zhangyang Wang, and Beidi Chen. H2o: Heavy- hitter oracle for efficient generative inference of large language models. In shient Confer- ence on Neural Information Processing Systems, 2023. httptttss//ppeneeeeeete wwweeeiieew.... net. Forum? d=
Dawei Zhu, Nan Yang. Liang Wang, Yifan Song. Wenhho Wu, Furu Wei, and Sujian Li. POSE: Efficient context window extension of LLMs via positional skip-wise training. In The Twelfth International Conference on Learning Representations, 2024. URL
set//oorum2id=32lloxuk
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. LongloRA: Efficient fine-tuning of long-context large language models. In The Twelfth International Confer- ence on Learning Representations, 2024b. TRLhttps::/ http:::///ppereviewwe net forum? id
6PmJoRfdaK
Weizhi Wang. Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. Augmenting language models with long-term memory. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL htttps://openree http:/ee//eeeee.eee fere/foorm?iie
Br BryMF PQ416
Maor Ivgi. Uni Shaham, and Jonathan Berant. Efficient long-text understanding with short-text models. Transactions of the Association for Computational Linguistics, 11:284-299, 2023. doi 10.1162/tacl a 00547
Samuel J Gershman, Christopher D Moore, Michael Toodd, Kenneth A Norman, and Per B Sederberg. 2012.
The successor representation and temporal context. Neural Computation,
Marcus K Benna and Stefano Fusi. Place cells may simply be memory cells: Memory compression leads to spatial tuning and history dependence. Proceedings of the National Academy of Sciences, 118(51):e201442))) 18, 2021.
Charles Blundell, Benigno Uria, Alexander Pritzel, Yazhe Li, Avraham Ruderman, Joel Z. Leibo, Jack Rae, Daan Wierstra, and Demis Hassabis. Model-free episodic control. arXiv preprint arXiv:1606.04444 2016.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adrià Puigdomênech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pages 2827-2836. PMLR, 06-11 Aug 2017.
Julian Coda-Forno, Changmin Yu, Qinghai Guo, Zafeirios Fountas, and Neil Burgess. Leveraging episodic memory to improve world models for reinforcement learning. In Memory in Artificial and Real Intelligence (MemARI) Workshop at 36th Conference Neural Information Processing Systems (NeurIPS 2022), 2024.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barrd Barwinska, et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114 (13):3521-352) -5226,,,,356,, 2017.
David Lopez-Paz and Marc' Aurelio Ranzato. Gradient episodic memory for continual learning. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett, editors, Advances in Neural Information Processing Systems, volume 30, Curran Asso- ciates, Inc., 2017. httpstttsts///p/eeeeeeiiii wwwrpps...........e............. httppppppccppppp..pppp felrffieesppppp /ess//pprr/////
Arslan Chaudhry, Mare' Aurelio Ranzato, Mareus Rohrbach, and Mohamed Elhoseiny. Efficient lifelong learning with a-GEM. In International Conference on Learning Representations, 2019. URL httprnreeeeeeeeeeo/ooreee22
14

---

Published as a conference paper at ICLR 2025
Pietro Buzzega, Matteo Boschini, Angelo Porrello, Davide Abati, and SIMONE CALDER- ARA. Dark experience for general continual learning: a strong, simple baseline. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 15920-15930. Curran Asso- ciates, Inc., 2020.
neurips cc/paper files/paper Papere.ppppppp...
Ameya Prabhu, Philip H. S. Torr, and Puneet K. Dokania. Gdumb: A simple approach that questions our progress in continual learning. In Andrea Vedaldi, Horst Bischof, Thomas Brox, and Jan- Michael Frahm, editors, Computer Vision ECCV 2020, pages 524-540, Cham, 2020. Springer International Publishing.
Payel Das, Subhajit Chaudhury, Elliot Nelson, Igor Melnyk, Sarath Swaminathan, Sihui Dai, Aurélie Lozano, Georgios Kollias, Vijil Chenthamarakshan, Soham Dan, et al. Larimar: Large language models with episodic memory control. arXiv preprint arXiv:2403.11990 2024.
Eleanor Spens and Neil Burgess. generative model of memory construction and consolidation. Nature Human Behaviour, pages 1-18, 2024.
Qihong Lu, Uri Hasson, and Kenneth A Norman. A neural network model of when to retrieve and encode episodic memories. elife, 1 11:e74455, 2022.
Maxine T Sherman, Zafeirios Fountas, Anil K Seth, and Warrick Roseboom. by-ill predictions of subjective time from human brain activity. PLOS Computational Biology, 18(7)se1010223, 2022.
Alexey Zakharov, Qinghai Guo, and Zafeirios Fountas. Variational predictive routing with nested subjective timescales. In International Conference on Learning Representations, 2022a. URL htttp:://ppenrew: ew net www..eeefformmiicm f=
Alexey Zakharov, Qinghai Guo, and Zafeirios Fountas, Long-horizon horizon video prediction using a dynamic latent hierarchy. arXiv preprint arXiv:22/2.14376 2022b.
Alexey Zakharov, Matthew Crosby, and Zafeirios Fountas. Episodic memory for subjective-timeseat models. In ICML 202/ Workshop on Unsupervised Reinforcement Learning, 2021. URL https 7/openreview
Nelson Cowan. The magical number 4 in short-term memory: A reconsideration of mental storage capacity. Behavioral and brain sciences, 24(1):87-114, 2001.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference Learning Representations, 2024b. B fhtpps: http///perrreee
net fenst/forunt ft/./formmmmiio diw.ummmiiiiGG s551 zVF
Chi Han, Qifan Wang, Hao Peng. Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. LM-infinite: Zero-shot extreme length generalization for large language models. In Kevin Duh, Helena Gomez, and Steven Bethard, editors, Proceedings of the 2024 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers). pages 3991-4008 Mexico City, Mexico, June 2024b. Association for Computational Linguistics. URL httpt:ttpp.............tt //aps://aclaaaae antho hrrpchhoogy...o ogy. wwt.////222222 naac ong. 222
Pierre Miasnikof, Alexander Shestopaloff, Anthony Bonner, and Yuri Lawryshyn. A Statistical Performance Analysis of Graph Clustering Algorithms, pages 170-184. 05 2018. ISBN 978-3- 319-92870-8. doi 11000.00779778-33319-92887
Mark EJ Newman and Michelle Girvan. Finding and evaluating community structure in networks. Physical review E, 69(2) 2004.
Matthijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. The faiss library. 2024.
Chankyu Lee, Rajarshi Roy, Mengyao Xu, Jonathan Raiman, Mohammad Shoeybi, Bryan Catanzaro, and Wei Ping. Ny-embed: Improved techniques for training llms as generalist embedding models. arXiv preprint ar0iv: 17428, 2024.
15

---

Published as a conference paper at ICLR 2025
Zhuowan Li. Cheng Li, Mingyang Zhang, Qiaozhu Mei, and Michael Bendersky. Retrieval augmented generation or long-context llms? a comprehensive study and hybrid approach, 2024b. URL
org/ abs 2407.16833
Sebastian Michelmann, Amy Price, Bobbi Aubrey, Camilla Strauss, Werner K Doyle, Daniel Friedman, Patricia C Dugan, Orrin Devinsky, Sasha Devore, Adeen Flinker, et al. Moment-by- moment tracking of naturalistic learning and its underlying hippocampo-cortical interactions. Nature communications, 12(1):5394, 2021.
Olga Lositsky, Janice Chen, Daniel Toker, Christopher Honey, Michael Shvartsman, Jordan Poppenk, Uri Hasson, and Kenneth A Norman. Neural pattern change during encoding of narrative predicts retrospective duration estimates. elife, 1:e1 2016.
Alberto Mariola, Zafeirios Fountas, Lionel Barnett, and Warrick Roseboom. Event segmentation in continuous, naturalistic videos from model- data-driven, and human perspectives. 2022.
Sebastian Michelmann, Manoj Kumar, Kenneth Norman, and Mariya Toneva. Large language models segment narrative events similarly to humans. arXiv preprint 2rX11:2301.10297, 2023b.
Alan Baddeley. Working memory: looking back and looking forward. Nature reviews neuroscience, 4(10):829-833) -9988 2003.
K Anders Ericsson and Walter Kintsch. Long-term working memory. Psychological review, review, 102(2): 211, 1995.
Santo Fortunato. Community detection in graphs. Physics reports, 486(3-5):75-174, 2010.
Zhao Yang, René Algesheimer, and Claudio J Tessone. comparative analysis of community detection algorithms on artificial networks. Scientific reports, 6(1):30750, 2016.
Ryan Prescott Adams and David JC MacKay. Bayesian online changepoint detection. arXiv preprint 1xXiv:0710.37444 2007.
Nils Reimers.
all-mpnet-base-v2. all mpnet-base httpp:://ncgiiggfees co http:/eeteeeeeeeaaaoooeeeeeeeee http:////aggggaaaeeeeoooooooooor
Sentence-transformmrs::
Hugging Face Model Hub,
2022. Available from:
Niklas Muennighoff, Nouamane Tazi, Loie Magne, and Nils Reimers. Mteb: Massive text embedding benchmark. arXiv preprint arXiv:2210.07311 2022.
Zhuowan Li, Cheng Li, Mingyang Zhang, Qiaozhu Mei, and Michael Bendersky. Retrieval augmented generation or long-context Ilms? a comprehensive study and hybrid approach. arXiv preprint arXiv:2407.16833, 2024c.
Rolf Jagerman, Honglei Zhuang, Zhen Qin, Xuanhui Wang, and Michael Bendersky, Query expansion by prompting large language models, 2023. URL WtLLttts//aaaxvoooo wttp://wwwrgrraa//////a///////// http//bab222030000
Yunxiao Shi, Xing Zi, Zijing Shi, Haimin Zhang, Qiang Wu, and Min Xu. Eragent: Enhancing retrieval-augmennt language models with improved accuracy, efficiency, and personalization, 2024. httpt::/p http:aawwww.rggbsa/aa/aaaaaasa/a
06683
Sagnik Majumder, Chinmoy Samant, and Greg Durrett. Model agnostic answer reranking system for adversarial question answering. CoRR, abs/2102.03016, 2021. URL /RL.attp:://axxv.oog//
abs/ / 22//22122.0 03016
Victor M. Panaretos and Yoav Zemel. Statistical aspects of wasserstein distances. nual Review of Statistics and Its Application, 6(Volume 6. 2019):405-43,, 2019.
An-
ISSN 2326-831X.
doi http://wwwaaaaaaaaoorerrrrrrr/o/
URL https: wwwp///wwww.k
annurev Swuueev-taatstics 000.--s--33011-0033333 0000033811LL11100337773113333333 http://wwwwwr///ooeeeeeeeeeeeeee
Tianzuu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. Differential transformer. arXiv preprint arXiv:2410.05258. 2024.
16

---

Published as a conference paper at ICLR 2025
Peng Wang. Shuai Bai, Sinan Tan, Shijie Wang, Zhihao Fan, Jinze Bai, Keqin Chen, Xuejing Liu, Jialin Wang. Wenbin Ge, Yang Fan, Kai Dang, Mengfei Du, Xuancheng Ren, Rui Men, Dayiheng Liu, Chang Zhou, Jingren Zhou, and Junyang Lin. Qwen2-vl: Enhancing vision-language model's perception of the world at any resolution, 2024. URL htttss://arxxtt www.aawwww......aa.aa.aooaooosre wttp://..rr.ab////gg//0
12191
Pravesh Agrawal, Szymon Antoniak, Emma Bou Hanna, Baptiste Bout, Devendra Chaplot, Jessica Chudnovsky. Diogo Costa, Baudouin De Monicault, Saurabh Garg, Theophile Gervet. Soham Ghosh, Amélie Héliou, Paul Jacob, Albert Q. Jiang. Kartik Khandelwal, Timothée Lacroix, Guillaume Lample, Diego Las Casas, Thibaut Lavril, Teven Le Scao, Andy Lo, William Marshall, Louis Martin, Arthur Mensch, Pavankumar Muddireddy, Valera Nemychnikova, Marie Pellat, Patrick Von Platen, Nikhil Raghuraman, Baptiste Rozière, Alexandre Sablayrolles, Lucile Saulnier, Romain Sauvestre, Wendy Shang. Roman Soletskyi, Lawrence Stewart, Pierre Stock, Joachim Studnia, Sandeep Subramanian. Sagar Vaze. Thomas Wang. and Sophia Yang. Pixtral 12b, 2024. URL /ttpss:/aazxii org. abs/ 2410. 07073
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, Victor Greiff, David Kreil, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need, 2021. URL htlltttss///xxiv..oom /ttpe//////...tttt///////////o//
2008 02217
Vincent D Blondel, Joup-Loup Guillaume, Renaud Lambiotte, and Etienne Lefebwre. Fast unfolding of communities in large networks. Journal of statistical mechanics: theory and experiment, 2008 (10):P10008, 2008
Nicholas Franklin, Kenneth Norman, Charan Ranganath, Jeffrey M Zacks, and Samuel Gershman. Structured event memory: A neuro-symbolic model of event cognition. Psychological review, 127(3):327, 2020.
Zhuxi Jiang, Yin Zheng, Huachun Tan, Bangsheng Tang, and Hanning Zhou. Variational deep embed- ding: An unsupervised and generative approach to clustering. arXiv preprint arXiv: 161 1.05148, 2016.
John Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(4) 1982.
Alex Graves. Neural turing machines. arXiv preprint
2014. 2014.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626): 471-476 2016.
Hubert Ramsauer, Bernhard Schäff, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint vrXiv:2008.022171 2020.
Greg Wayne, Chia-Chun Hung, David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska- Barwinska, Jack Rae, Piotr Mirowski, Joel Leibo, Adam Santoro, et al. Unsupervised predictive memory in goal-directed agent. arXiv preprint arXiv: 1803. 10760, 2018.
Randall C O' O'Reilly and Kenneth A Norman. Hippocampal and neocortical contributions memory: Advances in the complementary learning systems framework. Trends in cognitive sciences, 6(12): 505-510, 2002.
Zhitao Ying, Dylan Bourgeois, Jiaxuan You, Marinka Zitnik, and Jure Leskovec. Gnnexplainer: Generating explanations for graph neural networks. Advances in neural information processing systems, 32, 2019.
Zhitao Ying, Jiaxuan You, Christopher Morris, Xiang Ren, Will Hamilton, and Jure Leskovec. Hier- archical graph representation learning with differentiable pooling. Advances in neural information processing systems, 31, 2018.
17

---

Published as a conference paper at ICLR 2025
Alison R Preston and Howard Eichenbaum. Interplay of hippocampus and prefrontal cortex in memory. Current biology, 23(17):R744-R73) 2013.
Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork variational autoencoders. Advances in Neural Information Processing Systems, 34:29246-29257, 2021.
Danijar Hafner, Kuang-Huce Lee, Ian Fischer, and Pieter Abbeel. Deep hierarchical planning from pixels. Advances in Neural Information Processing Systems, 35:26091-26104, 2022.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with gpus. IEEE Transactions on Big Data, 7(3):535-541) 2019.
18

---

List of Appendices
A Analytical results
19
A.Coompaiioon with V-rettiieaal based LLMs
19
A.2Comparison with RAG
23
Human data
25
B. Anaalysis
25
B.2 Further results
26
Complexity
27
C. Boundary Refinement Complexity Analysis
27
C.2Attention Complexity Analysis
................................
27
C.3 ImIpleementaioon and use of hardware resources
28
D Further Ablations
30
................................
D. H.peerrr-rrraete Selection and Tuning
30
D.2Surprise, refinement and contiguity
................................
31
D.3Retrieved Tokens and Context Length
33
Further Discussion points
35
E. Detailed Analysis of EM-LLM's Connection to Human Episodic Memory
35
H.2 Architecture contributions EM-LLM
36
E.3 Why is the Refinement Algorithm Elfective?
37
E.4 Feasibility of End-to-End Neural Implementations
37
E.5 Future Extensions Inspired by Human Memory Systems
38
Proofs
38
A ANALYTICAL RESULTS
A.I COMPARISON WITH KV-RETRIEVAL-DASED LLMS
While the vast majority of our results in this section show a statistically significant improvement on InfLLM at the benchmark level (p 0.05 using two-tailed z-test, except LongBench with 0.23), it should be noted this isn't the case in the majority of individual tasks.
Phi-3.5 with However, given the consistency and frequency of improvements across large number of such tasks, along with the benchmark-level significance of such improvements. we consider the lower, task-level significance to be largely due to the sample size individual tasks rather than chance, and believe still reasonable and justified claim overall improvement InfLLM. Moreover, including individual task results supports transparency and allows for future works to make more granular comparisons and use of such results.
19

---

Published as a conference paper at ICLR 2025
EM-LLM
Task Type
Task
InfLLM
Max Imp.
SM
S+C
SM+C
Single-doc QA
NarrativeAA
22.12
1.49
21.77
21.13
21.10
22.45
Qasper
29.33
0.17%
29.07
29.38
29.16
28.68
MultiFieldQA
47.42
1.46 %
48.11
47.39
47.72
47.62
Multi-doc QA
HotpotQA
36.56
10.15% %
39.40
39.01
40.27
38,90
2WikiMQA
22.31
5.20 %
23.46
22.75
23.47
23.46
Musique
17.68
6.17%
17.97
17.82
17.98
18.77
Summarisation
GovReport
31.03
1.90 %
31.40
31.62
31.10
31.43
QMSum
23.49
2.13%
23.99
23.20
23.48
23.47
MultiNews
26.70
-0.30%
26.55
26.54
26.58
26.62
Few shot
TREC
69.00
2.90%
71.00
70.00
70.50
70.50
TriviaQA
86.67
1.10%
86.58
87.62
87.52
87.47
SAMSum
42.52
0.45%
42.71
42.13
42.34
42.48
Retrieval
PassageRetrieval
64.00
32.69%
82.67
78.92
84.92
84.08
Code
LCCC
56.67
0.64%
55.03
57.03
54.90
56.79
RepoBench-P
52.97
1.34%
50.49
53.68
51.06
52.86
Avg. score:
41.90
4.50
43.35
43.22
43.47
43.71
Code Multiple choice
Code.
29.44
0.88 %
29.70
28.43
28.68
28.17
En.MC
43,23
-1.02%
41.48
40.61
42.79
42.79
Retrieval
Math.Find
26.57
5.38% %
28.00
27.43
27.71
27.14
Retrieve.KY
95.60
3,56
92.20
97.20
97.60
99.00
Retrieve. Retrieve.PassKey
100.00
0.00%
100.00
100.0
100.00
100.00
Porrive.Nuubber
99.83
0.00%
99.83
99.83
99.83
99.83
Avg. score:
65.78
1.47
65.20
65.58
66.10
66.16
Table 3: EM-LLM performance on LongBench and oo-Bench (respectively) compared to our baseline, as the base L.M and 4K+2K context. S: surprise threshold, SM: surprise threshold + refinement with modularity, S+C: surprise threshold + contiguity buffer, SM+C: surprise threshold + refinement with modularity contiguity buffer. Max Imp.: Maximum relative improvement over InfLLM across all EM-LLM variants.
InfLLMM, with
20

---

Published as a conference paper at ICLR 2025
EM-LLM
Task
InfLLM
Max Imp.
SM
S+C
SM+C
NarrativeQA
22.64
8.92%
24.47
22.50
24.66
23.03
Qasper
43.70
3.25%
44.35
44.95
45.07
45.12
MultiFieldQA
49.03
0.47%
49.11
48.79
49.26
48.36
HotpotQA
49.04
0.69%
48.97
49.19
48.74
49.38
2WikiMQA
35.61
8.26%
38.44
38.08
38,55
38.08
Musique
26.06
-1.46%
25.68
25.19
24.64
23.92
GovReport
30.76
1.11%
31.10
30.85
30,96
30.86
QMSum
22.70
0.31%
22.63
22.77
22.62
22.58
MultiNews
27.57
-0.91%
27.32
27.28
27.30
27.29
TREC
73.50
0.00%
73.50
73.50
73.50
73.50
TriviaQA
90.91
0.00%
90.91
90.91
90.91
90.91
SAMSum
42.43
1.98%
43.24
43.27
42.91
42.84
PassageRetrieval
84.00
4.17%
87.50
86.00
86.00
85.00
LCC
59.88
0.94%
58.49
60.44
58.55
60.41
RepoBench-P
46.48
-3.44%
42.13
44.88
42.26
44.68
Avg. score:
46.95
1.62%
47.19
47.24
47.06
47.06
Code. Code.Debug
30.46
5.81% %
31.73
30.20
32.23
31.73
Math.Find
23.70
27.668%
16.86
16.57
16.29
17.14
Retrieve.KV
5.00
4.00%
4.20
4.80
5.20
En.MC
43.70
-3.07%
40.55
42.36
39.74
40.61
Retrieve. Payssky
100.00
0.00%
100.00
100.00
100.00
100.00
Cetrieve.Number
99.00
0.67% %
99.49
99.49
99.66
99.49
Avg. score:
50.31
-3.38%
48.81
48.90
48.85
49.00
Table 4: EM-LLM performance LongBench and oc- Bench (respectively) compared to our baseline, InfLLM, with LLaMA-3-aR ----ruuttctt as the base LLM and 4K+4K context. Abbreviations as before.
EM-LLM
Task
InfLLM
Max Imp.
SM
S+C
SM+C
NarrativeQA
26.64
2.25%
26.05
26.05
27.24
25.98
Qasper
44.95
1.27%
44.41
45.52
44.71
45.41
MultiFielQQA
52.56
-0.08%
52.52
52.07
52.36
52.48
HotpotQA
52.96
2.00%
54.02
52.49
53.37
53.90
2WikiMQA
45.04
4.57%
45.72
45.60
46.61
47.10
Musique
23.98
7.76%
25.37
25.84
24.60
24.73
GovReport
34.96
0.57%
35.04
35.16
34,94
34.81
QMSum
24.36
0.94%
24.31
24.55
24.31
MultiNews
27.78
0.14%
27.76
27.79
27.82
27.77
TREC
71.00
0.70%
71.50
71.50
71.50
71.00
TriviaQA
92.44
0.00%
92.34
92.24
92.43
92.44
SAMSum
43.68
0.41%
43.31
43.65
43.63
43.86
PassageRetrieval
97.00
2.58%
98.50
98.00
97.50
LCC
65.82
2.48%
67.45
65.69
65.74
65.74
RepoBench-P
62.56
2.83%
64.33
62.54
62.18
61.87
Avg. score:
51.05
1.89%
51.58
51.28
51,30
51.28
Code. Code.Debug
22.59
0.00%
22.59
22.59
22.59
22.59
Math.Find
33.71
6.79%
36.00
34.00
35.43
27.71
Retrieve.KY
81.00
19.51 %
96.80
90.20
95.40
95.20
En.MC
46.72
1.88% %
44.54
47.60
46.72
45.85
Sestieve.PassKer
100.00
0.00%
100.00
100.00
100.00
100.00
Cetrieve.Number
100.00
0.00%
100.00
100.00
100.00
100.00
Avg. score:
64.00
4.70% %
66.66
65.73
66.69
65.23
Table 5: EM-LLM performance LongBench and co- so-Bench (respectively) compared to our baseline,
InfLLM, with LLaMA-3.1-8BB- I-8B-Instrucct as the base LLM and 4KK context. Abbreviations as before.
21

---

Published as a conference paper at ICLR 2025
EM-LLM
Task
InfLLM
Max Imp.
SM
S+C
SM+C
NarrativeQA
14.82
14.04%
15.78
16.90
16.02
16.66
Qasper
28.71
5.29%
28.25
28.46
29.10
30.23
MultiFieldQA
41.54
5.66%
43.48
43.89
42.57
42.57
HAtpotQA
32.64
10.45
36.05
34.37
33.53
31.98
2WikiMQA
27.08
10.16 %
28.74
28.05
29.83
27.49
Musique
15.05
29.70 %
16.53
16.76
19.52
16.49
GovReport
28.96
2.97 %
29.41
29.59
29.82
29.62
QMSum
21.64
3.00%
22.29
22.06
22.07
MultiNews
26.32
-0.49%
26.07
26.16
25.89
26.19
TREC
67.00
3.73%
68.50
69.50
68.50
68.00
TriviaQA
83.71
1.73%
83.59
84.16
85.16
85.09
SAMSum
7.83
8.30%
8.43
8.48
7.28
8.06
PassageRetrieval
7.50
40.00%
10.00
10.50
LCC
60.33
-0.22%
59.99
60.13
60.20
60.01
RepoBench-P
53.70
0.37 %
53.90
53.59
53.44
53.27
Avg. score:
34.46
8.98%
35.40
35.43
35.46
35.22
Table 6: EM-LLM performance on LongBench compared to our baseline, InILLM, with Phi-3-Mini- 4K-Instruct the base LLM and IK+KK context. Abbreviations before.
EM-LLM
Task
InfLLM
Max Imp.
S
SM
S+C
SM+C
NarrativeQA
17.82
8.98%
19.42
16.63
17.55
17.12
Qasper
31.44
3.34%
32.38
31.71
31.43
32.49
MultiFieldQA
45.80
-2.36%
43.58
44.72
44.28
44.66
HotpotQA
41.33
11.35% %
46.02
44.78
44.43
44.89
2WikiMQA
27.74
8.51
29.68
30.10
29.27
Musique
16.39
21.23%
19.87
19.26
18.70
19.41
GovReport
26.37
3.22 %
27.05
27.22
26.76
27.13
QMSum
21.19
4.29 %
21.94
22.10
21.79
21.79
MultiNews
24.23
0.70%
24.39
24.29
24.29
24.40
TREC
67.50
2.22%
67.50
67.50
68.00
TriviaQA
84.66
-0.99%
83.82
83.82
83.82
83.49
SAMSum
16.62
-0.54%
15.30
14.49
16.53
15.25
PassageRetrieval
11.50
17.39%
13.00
13.50
13.00
13.50
LCC
38.38
-3.20%
36.79
37.15
36.90
37.02
RepoBench-P
42.30
1.75%
42.16
43.04
41.91
41.65
Avg. score:
34.22
5.06 %
34,86
34.69
34.60
34.67
Table 7: EM-LLM performance on LongBench compared to our baseline, InfLLM, with Phi-3.5- mini-Instruct as the base LLM and 1K3KK context. Abbreviations as before.
22

---

Published as a conference paper at ICLR 2025
A.2 COMPARISON WITH RAG
EXPERIMENT DETAILS
In our experiments with Rerrieval-Augmented Generation (RAG) baselines, we implemented a standard RAG pipeline consisting of a retriever and a downstream LLM. For each example in benchmark task, the example context is split into chunks of words each and encoded using the retriever's s embedding model into a vector database. A similarity lookup into the vector database is used to retrieve the top k most relevant chunks, which are then fed to the downstream LLM alongside the query and task description.
We conducted experiments using two retriever models NV-Embed-v2 (Lee al., 2024 and all- mpnet-base-v? (Reimers) 2022). NV-Embed-v2 is a SOTA LLM-basee retriever that uses that, as of September 2024, ranks first the Massive Text Embedding Benchmark (MTEB) Leaderboard (Muennighoff et al. 2022). It is a fine-tuned Mistral-7BvO, model with an emhedding dimension of 4096, trained using contrastive instruction-tuning on both retrieval and non-retrieval datasets. all-mppet-base-v2 is smaller 110M parameter model with an embedding size of 768, built on the BERn-base-uncased architecture, trained using contrastive learning on a dataset of over 1 billion sentence pairs. For each embedding model, we ran experiments using LLaMa-3-88 and LLaMa-3.1- 8B as the downstream LLM.
For most experiments, we set chunk size to I 300 and k 5, following the protocol of Li et al. (2024c). As a comparative baseline, we also trialled chunking according to surprise, the results of which are shown in Table (RAG-S). In this experiment, the context was first segmented into blocks by the EM-LLM (S) model, each of which then encoded by the retriever's embedding model. Context was dynamically retrieved fill buffer of 1500 words, for fair comparison with fixed-size chunking. The RAG-S variant underperformed the fixed-size implementation we hypothesise that this was due to the retrieved context in RAG-S consisting of highly disjointed information, due to small chunk sizes and high k. Incorporating contiguity into the retrieval mechanism may be sufficient to close this performance gap.
LIMITATIONS OF RAG
RAG requires the use of additional modules alongside the downstream LLM during the generation process, meaning that the quality of the generated output depends on the representational capacity of these modules in addition to the capability of the LLM. For example, the use of retriever model that has far fewer parameters than the downstream LLM can limit the LLM's ability to generate the most accurate or contextually appropriate responses, as shown by the gap in performance between the two RAG pipelines on LongBench shown in Table [ In this case, whilst the downstream LLM may be capable of high performance on given task, the retriever may be expressive enough to provide the relevant context needed to solve said task. Additional pre/post-retrievall techniques such as query expansion (Jagerman et al. (2023). ). knowledge filtering (Shi et al. (024g) or answer reranking (Majumder et al. (2021)), may help to bridge potential performance bottlenecks, but these involve further increasing the complexity of the generation pipeline. In contrast, EM-LLM outperforms both RAG models whilst only requiring the use of a single LLM across the entire generation stage, removing the issue of performance bottlenecks seen in RAG pipelines. Furthermore, EM-LLM is general purpose method that can be applied to practically any existing transformer LLM method was implemented using a general purpose patch to attention layer modules that provided compatibility with the Huggingface Transformers Library.
23

---

Published as a conference paper at ICLR 2025
RAG
Task
EM-LLM
all-mpnet-base-22
NV-Embed-v2
NarrativeQA
17.31
21.39
22.50
Qasper
40.66
41.01
44,95
MultiFielQQA
45.16
50.47
48.79
HotpotQA
43.32
53.64
49.19
2WikiMQA
41.48
40.41
38.08
Musique
26.78
31.56
25.19
GovReport
27.87
29.26
30.85
QMSum
22.44
23.15
22.77
MultiNews
26.04
27.48
27.28
TREC
4.50
65.00
73.50
TriviaQA
78.98
63.75
90,91
SAMSum
9.00
32.85
43.27
PassageRetrieval
54.00
54.50
86.00
LCC
10.76
19.88
60.44
RepoBench-P
13.02
34.77
44.88
Avg. score:
30.75
39.27
47.24
Table 8: Comparison of RAG with two different retrievers vs. EM-LLM on the LongBench dataset. All methods use LLaMs-3-8B-Instruutt for generation.
Task
RAG-S
RAG
FC
EM-LLM
NarrativeQA
12.10
22.54
29.14
26.05
Qasper
36.41
45.45
45.34
44.41
MultiFielQQA
44.08
51.67
54.98
52.52
HotpotQA
41.56
55.93
54.01
54.02
2WikiMQA
28.84
42.93
45.95
45.72
Musique
19.04
30.90
33.52
25.37
GovReport
18.12
29.91
34.49
35.04
QMSum
19.22
24.97
25.14
24.31
MulúNews
26.21
26.77
27.00
27.76
TREC
2.5
22.50
4,50
71.50
TriviaQA
88.26
88.11
89.07
92.34
SAMSum
8.09
7.56
8.68
43.31
PassageRetrieval
16.0
65.50
100.00
99.50
LCC
11.02
13.16
19.30
67.45
RepoBench-P
17.39
18.66
18.33
64.33
Avg. score:
25.89
36.44
39.30
51,58
Code.Debug
22.59
21.70
22.59
Math.Find
35.43
26.29
36.00
Retrieve.Ky
31.80
92.60
96.80
En.MC
64.19
58.07
44.54
Retrieve, Passkvv
100.00
100.00
100.00
Retrieve.Number
99.83
99.32
100.00
Avg. score:
58.97
66.33
66.66
Table 9: EM-LLMg (4K+4K) RAG (NV-Embed-v2 retriever) vs. full-context, with LLaMa-3.11
8B as the base LLM, evaluated on LongBench and oo-Bench. The comparison also includes RAG-S, which is the same RAG retriever but with the surprise-based segmentation used in EM-LLM results.
24

---

Published as a conference paper at ICLR 2025
Ratio of Unique Memory Blocks Retrieved across Transformr Layers and Time
80
60-
15
18
32
Trarsformrr layer
100
20-
60
16-
pallo
12
25600
51200
76800
102400
128000
153600
179200
204B00
230400
256000
Token Index
Figure 5: The ratio of blocks retrieved by a layer which were not retrieved by any other layer for the same processed chunk, versus the total number of retrieved blocks by that layer. This is measured using EM-LLM, with Mistral-7B on a single example of oc- co-Bench's Longbook. Choice. Eng task, with over 500 chunks of 512 tokens. RAG methods, this ratio would always be zero, as retrieved blocks are used by all layers concurrently.
B
HUMAN DATA
B. ANALYSIS
The human data released part of Kumar al. (2023) used Gaussian smoothing on the average signal across participants to define a probability distribution of likely event boundary positions with respect timestamps in the podcast. In order to calculate our similarity metrics, as shown in Fig. 4A. we need to express this data in terms of discrete event positions with respect to tokens in the transcript. For fair comparison, we therefore identified most likely positions in the distribution as our initial curprise-based event segmentation had identified the transcript. the same process used by Kumar al. (2023), we then used their provided word onset times to translate these timestamps to token positions, allowing calculate our similarity metrics.
positions by selecting as many of the
In Fig. 44, we use Wasserstein distance in order to compare the relative positions of boundaries between human annotations and those found by our own methods. Wasserstein distance is versatile metric used to compare two probability distributions (Panaretos and Zemel 2019). We used such metric to better capture the uncertainty present in the human data, and found it to give more meaningful results than standard correlation or discrete distance metrics, which showed very little differences between methods. In order to calculate such metric, we therefore need to convert our own discrete boundary positions to a distribution across token positions. We did so by defining a Mixture of Gaussians (MoG), with each Gaussian corresponding to single position. Note that, for fair comparison with human data, we apply the same process to the discrete version of the human-annotate positions described above, and use this for comparison.
25

---

Published as a conference paper at ICLR 2025
A
Monkey
B
2.5
1.81
2:40
Wassersten in distance human reports (x1
Modulanity
1.5
PC
Random: 1.66
Random: 2,06
Humana
inra-liter-similary
1.44-
2.03-
0.56
1,02-
0.46
1.00
1.01
0.33
0.25
0.25
0.18
25
30
15
20 30
P51 5M
AC
5C
PM
5M
PC
52
LLM layer
LLM
Figure 6: Comparison of human event segmentation with different computational segmentation audio datasets. (A) Difference in metrics for the cohesion and
methods in two separation KV cache of LLaMA2 attention heads. The graphs report the difference of each method with the corresponding random segmentation. (B) Distance between human reports and different methods. In both sets of results, fixed methods (F, FM, FC) perform worse than their surprise-base counterparts (S, SM, SC) with InfLLM's method (F) performing worse than random.
Monkay
Monkey
Tunnel podcast
Wes Passserstteil distancett human regorts De
Modularity
1.6
1.04
2.4
1.52
2.00
layer
Figure 7: Comparison of human event segmentation with different computational segmentation methods using Mistral-7B. Plots include abbreviations like before.
B.2 FURTHER RESULTS
Monkety podcast
Tunnel podcast
Monkey padcast
Tunnel podcast
Wesserste distanncet human reports (X10)
Modularit
1.6.
1.4
Intre-fnn intre-inntersing wnteerimmillrrrri
0.64 0.11
1.22
031 0.20
LLM
Figure 8: Comparison of human event segmentation with different computational segmentation methods using LLaMA-3-88 Plots include abbreviations like before.
26

---

Published as a conference paper at ICLR 2025
C
COMPLEXITY
C.I BOUNDARY REFINEMENT COMPLEXITY ANALYSIS
Here, we provide detailed analysis of the computational complexity of our Algorithm|] T focusing on the boundary refinement step and the calculation of modularity and conductance metrics. Here we describe scaling complexity with example context length n, chunk size and k events.
The metric functions (modularity or conductance) are computed at the level of individual memory units but all rely on the same adjacency matrix for that chunk. The complexity of calculating the adjacency matrix, as defined in Eq. 2 is O(m²). As the complexity for both metrics is largely dominated by, or on the same order as the computation of this term, the resulting complexity for the calculation of the metric function is O(m²).
Metric Function Computation
Boundary Refinement Step The boundary refinement step involves finding the optimal position 3 between each pair of consecutive initial boundaries (a, .3 that optimizes the chosen metric function f. The iteration over initial boundaries has complexity (k)). For each updated boundary we change the community of one node at at time to the one on its right and re-evaluate the metric function with this change. Hence, for each boundary position, on average 2% nodes see a change in their community, and hence need to re-evaluate their contribution to the metric function. This results in complexity of O 2 ). for each position between a and 3. This step will therefore scale with average event size resulting in complexity (?) Therefore, the overall complexity of this step scales with (k))))) (
Overall Complexity The context is divided into a chunks, and hence, must compute this number of initial adjacency matrices with resulting complexity ( (?) O(m)). Adding in the complexity of the refinement updates, we get (num + =² In practice, our method detects k events such that n S m is always true, hence this is upper-bounded by O(nm).
C.2 ATTENTION COMPLEXITY ANALYSIS
Attention. Calculating a full attention matrix for a sequence length n has complexity O(?)). However, in the case of EM-LIM, we process such a context in smaller chunks of size m and rely on retrieval to approximate full attention matrix (also see Appendix F.T). Therefore we evaluate a attention matrices each with complexity Omm(n + n,,)) with ny the number of local tokens used and n, the number of tokens retrieved as part of the multi-stage attention process.
k-NN Retrieval. As previously mentioned, our approach relies on k-NN retrieval to avoid computing a full attention matrix. Such retrieval process scales with the maximum number of memory units saved. our case, as we enforce minimum block size 6, the maximum possible number of memory units to check for retrieval is n nn As we retrieve memory units for every processed chunk (past my
processed tokens) this is upper-bounded by O p2(n-n))
mub
Overall. Including the retrieval step, overall our approach performs its attention computation with complexity O(n(ny + n2,) + m(n-ni)). ). In practice, for long-contex, tasks, this scales much slower than full attention matrix, as visualized in Figures Jand 10 which demonstrate this with our settings and for values of n up to 100K and 10M respectively. In fact, our complexity negligible compared to full attention
mub
sequence length reaches the millions.
27

---

Published as a conference paper at ICLR 2025
Context Attention
EM-LLM Attention
EM-LLM Attention Retrieval
1e10
1e8
1.0
1010
6
Complexity (log scale)
0.8
Complexity
109
Complexity
4
0.6
109
0.4
2
107
0.0
10
0
100000
100000
100000
Sequence Length
Sequence Length
Sequence Length
Figure 9: A visualization of the scaling complexity of EM-LLM vs standard full-context approach (full attention matrix) as a function of sequence length (up to 100K tokens).
Context Attention
EM-LLM Attention
Attention Retrieval
lel4
1e10
1.0
Complexity (log scale)
1013
0.8
Complexity
10
Complexity
6
4
0.4
107
0.0
0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1e7
1e7
1e7
Sequence Length
Sequence Length
Sequence Length
Figure 10: A visualization of the scaling complexity of EM-LLM vs standard full-context approach (full attention matrix) as a function of sequence length (up to 10M tokens).
C.3 IMPLEMENTATION AND USE OF HARDWARE RESOURCES
Retrieval-based methods for handling long sequence lengths can be especially appealing to those who may not have access to the hardware resources necessary to run large, long-context models. In this spirit, we describe here our adjustments to our framework made to further lower the minimum hardware requirements accurately run inference on sequence lengths of 1M+ tokens. We note that our approach is very scalable in the sense that handling longer contexts only requires an increase in CPU memory, or even just disk storage.
All of our experiments (except for the full-context results for which we used an API). Additionally, each node had a minimum of 100GB of CPU memory. While all base models and methods used across our own experiments fit on a single GPU, such hardware is still quite limited when compared to the more advanced H100s commonly used in research nowadays. Furthermore, the memory overhead due processing and storing the KV cache and representative tokens for each memory unit for very long sequence lengths means that we have had to make some specific adjustments in order to ensure that we can efficiently run such experiments on older and limited hardware.
single nodes of 4 GPUs, each with 32GB of dedicated memory
C.3.1 WALL CLOCK TIME
As can be observed in Table 10, the largest increase to wall clock time in our framework is due to the similarity adjustment step. We would also like to note that such times increase steadily as sequence length increases, due the increasing number of representative tokens involved in the
28

---

Published as a conference paper at ICLR 2025
EM-LLM
InfLLM
S
S+C
SM
SM+C
Time
1.40s
1.57s
1.57s
2.27s
2.27s
Ratio
1.0
1.12
1.12
1.62
1.62
Table 10: Difference wall-clock time to process a single chunk of 512 tokens for various ablations of our framework as compared to Inf..LM. Measured using Mistral-7B (4K+2K) and averaged over the first 100 chunks (51.2K tokens) of a long sequence.
k-NN calculation used for block retrieval, regardless of which method is used, although this is only noticeable when sequence length reaches the millions (see Appendix C.2
C.3.2 MEMORY REQUIREMENTS
Seq. Length
KV Cache
Rep. Tokens (max)
20K
9.8GB 488.3GB
0.6GB
IM
30.5GB
Table 11: Memory requirements for components of our framework which scale with sequence length, for a model with 32 layers and 32 KV heads, assuming half float precision and batch size of 1. The number of representative tokens saved depends the number of blocks, hence show the maximum amount of memory required in the event the model segments the sequence into the maximum possible number of blocks.
As shown in Table 11 the KV cache can take up lot of memory for longer sequences. While the bulk is kept on CPU memory until it is needed, such a resource is quickly spent when running multiple instances in parallel on single node. Likewise, while the representative tokens for each block take up only a fraction of the memory compared to the KV cache, it can quickly become too large to keep GPU memory. In order to address both of these issues and facilitate the of our framework on lower-spec hardware, we have introduced various forms of offloading and model parallelisation.
CPU and Disk Offloading
InfLLM already provides details of their least recently used (LRU) KV cache management strategy used to efficiently offload the least-used memory units from GPU to CPU memory. We take inspiration from such a strategy and extend offload such LRU units to disk space when the number of available slots runs low. Allocating and overwriting the memory slots, as is also done in the GPU cache, prevents our process from hanging on the memory which has been freed, avoiding system fragmentation and allocation errors. Furthermore, the ability to offload to disk means that our framework only requires approximately 2GB of CPU memory per instance in order to run with little overhead, as long there is enough available disk space on the local machine.
use allocated memory slots in CPU memory and dynamically
Furthermore, we implemented the option to offload representative tokens CPU memory for very long sequences (IM+), although we note that this further increases the overhead seen in wall clock time during N-NN retrieval.
Model Parallelisation
For very long sequences (2.5M+) it may be more time-efficient to parallelise the model's layers across multiple GPUs in order to keep representative tokens on the same GPU and speed k-NN retrieval. We used Hugging Face's 's Accelerate to do this, along with some explicit casting of shared modules. Given a single node with set number of GPUs and sequence length below 1M tokens, it is faster to offload the representative tokens and run a separate instance on each GPU, than to parallelise on 2 GPUs.
29

---

Published as a conference paper at ICLR 2025
D FURTHER ABLATIONS
D.1 HYPER-PARAMETER SELECTION AND TUNING
Surprise Threshold Parameter
Equation| introduues the surprise threshold parameter y, which is responsible for the sensitivity of the threshold by scaling the standard deviation measured across the moving window. As such, with = 1, for a new token do be considered "surprising" it must achieve a surprise value greater than one standard deviation over the moving average. As mentioned in section 3.2 this ensures that the threshold adapts to contextual shifts therefore minimizing the need for manual tuning.
We initially explored our approach's sensitivity to 7 using Mistral on the LongBench benchmark. We evaluated the benchmark using surprise-only segmentation with y E 1.0, 1.5, 2.0, 2.5, 3.5, as visualized in Figure [2] Naturally, we noticed that an increase in 7 resulted in a decrease in the number of events detected, and resulting increase in the mean event size. In terms of overall performance, results suggest that a smaller 7, is generally best. We then explored such behavior across our segmentation methods, for which the overall behaviors were largely consistent with the surprise-only method. One particularly interesting observation is the fact the addition of our refinement algorithm (sM) seemed to show particularly high improvements over surprise-only larger values of although the best-performing value was still y = 1.
Following these initial observations, order to choose 7 for our experiments, we evaluated each model on the LongBench benchmark with 7 E 1,2,3 and surprise-only segmentation. We then selected the best best-performing value of ry for the rest of our experiments with each model. These were Y 2, 1, 1, for Mistral, LLaMa-3, LLaMa-3.1, Phi-3, Phi-3.5 respectively, showing consistent preference for 7 =
Retrieved vs. Contiguity Buffer Ratio
Section 3.4 introduces similarity buffer of k, events and a contiguity buffer of k. events which retrieves n. events either side of each event the similarity buffer. total, the number of retrieved events is k = k2 + k.- Note that we chose n = 1 for all experiments in order to balance contiguity with similarity. Moreover, Fig. AA&B suggest that the contiguity effect is most significant for handful of positions either side of the one recalled. This suggests that larger values of are likely to bring diminishing returns in terms of contiguity, while also reducing the size of the similarity buffer.
In practice, is expressed vary due to our dynamic segmentation methods. We therefore express the sizes of the similarity and contiguity buffers as a ratio k ig and dynamically fill them with events such that k2 = (1 k, kk and k = k, In order find the best contiguity ratio, we evaluated LongBench with Mistral and E 2 over values k, E 0.3, 0.5, 0.7, as illustrated in Fig. [3 To give a bit of intuition as to the meaning of these values, k, 0.3 will, on average, correspond to only including contiguous events for the top 25% of events in the similarity buffer. On the other hand, k, 0.7, on average, will correspond to including contiguous events for all events in the similarity buffer, but will include 50% less events in the latter. As mentioned in section 4.4, contiguity appears to have varying importance across tasks, with some performing best with the larger ke = 0.7, but most did best with the lower k, = 0.3, 0.5. Results also showed a clear preference for ? 1, which is consistent with our previous experiments. Overall, there was a slight preference for k, = 0.3, which we therefore selected for the rest of our experiments.
maximum number of tokens to ensure consistency as event sizes
E.
30

---

Published as a conference paper at ICLR 2025
D.2 SURPRISE, REHINEMENT AND CONTIGUITY
Best performing EM-LLM version across LLMs hyperparams
Best performing base LLM across all EM-LLM versions
40
1.0
2,0
14.
y=2.0
30
3.0
12-
y=
#experments
#experimetts
20
10
SURPRISAL
SIM. REF
CTG BUFF
HYBRID
Mistral
LL3
113.1
PH13
PH13.5
EM-LLM
Base
Figure 11: (left) Best performing version the EM-LLMM model across base language models (LMs) and hyper-parameters, including the parameter "y that controls the sensitivity of the segmentation threshold in Equation (T). (right) Best performing base LM across all experiments presented in this work.
31

---

Published as a conference paper at ICLR 2025
NarrativeQA
MultiReiDOA
HotpotOA
39.-
22-
30-
23
2WikiMQA
Musique
GovReport
QMSum
24.
22-
Mult/News
TREC
TriviaAA
SAMSum
48
5.0
10
83
PessageReitrise
56-
12-
Mean
Mean
Mean
Figure 12: Ablation study in LongBench with Mistrrl-7BB-Insttuutvv0012 Comparison of EM-LLM performance for different combinations of model features (represented by different colours) and different values of (the threshold's scaling factor). Model variants aligned the x-axis based on the average number block size that emerges for each case. The 7 values for each model variant shown in the first sub-plot. The corresponding InfLLM performance is also shown.
32

---

Published as a conference paper at ICLR 2025
NarrativagA
Qusper
MultiReldQA
Taak scome
23
24
228 255
250
218
20.0
2WikiMDA
Musique
GovReport
QMSum
24-
Thik score
175
223 SAMSum
Mull NNwws
TREC
TrivaQA
Contiguity rati:
Contiguityy rati: 0.5
Contiguity rato: 9.7
Thak Score
25.0
2500
200
2000
PassageRetrievel
LCC
RepoBench
Average scores of all tasks
Taak 5cone
29
27
22
34
26
26
28
30
30
34
30
26
20
34
38
16
18
27
22
24
26
28
Mean event size
Mean event size
Mean event size
Mean MV evintt size
Figure 13: Ablation study in LongBench with Mistral--B--Instrutttvve Comparison of EM-LLM performance for different ratios of the contiguity and similarity buffers (represented by different colours) and different values of y. Model variants aligned the x-axis based on the average number of block size that emerges for each case. The values for each model variant are shown in the first sub-plot. The corresponding InfL.LM performance is also shown.
D.3 RETRIEVED TOKENS AND CONTEXT LENGTH
In current experiments, have chosen buffer sizes align with related works (namely InfLLM) in order to make direct performance comparisons. Such values also keep buffer sizes shorter than the average number of tokens in the evaluated benchmarks ensure appropriate retrieval. In order to further explore variations in these parameters, we ran a small ablation study varying the size of the retrieved buffer for summarization tasks in LongBench (Table [2 and Fig. [4]. We have chosen such tasks as we believe they will be most likely to require more of the text content give accurate answer, and hence be most sensitive the number of retrieved tokens.
of
However, for the following reasons, we believe this provides limited information on such parameters. For such a study, we would choose a base LLM trained with a relatively large context window, such Mistral LLaMa 3.1 which support context lengths of up to 32K and 128K respectively, in order to ensure that the underlying model can support an adequate range of buffer sizes. As I LongBench may be considered relatively short compared these context windows (average number
33

---

Published as a conference paper at ICLR 2025
1K
2K
4K
6K
GovReport
QMSum
26
30
32
24
% of examples
Score
30
Score
22
20
28
% of
20
26
18
0K-4K
4K-8K
BK-12K
12K-16K
16k-20K
20K-2KK
24K-
0K-4K
4K-8K
BK-12
12K-16K
16k20K
20624K
24K+
MultiNews
SAMSum
26
20
% of examples
Scoee
40
Score
22
20
20
0K-4K
4K-BK
BK-12K
12K-16k
16k-20k
201-24K
24K+
0K-4K
4K-8K
12K-16k
16K-20K
201-24K
24k
BK-
Context length
Context length
Figure 14: Ablation of the size of the retrieved buffer EM-LLMs on Mistral with 4K local tokens on LongBench's 's summarization tasks. Presented as a function of context length.
of tokens per example: 12K 10K with Mistral's tokenizer), oc-Bench would be more appropriate (average number of tokens per example: > 100K). Unfortunately, evaluating larger buffer sizes (hence larger attention matrices) on the already-expensive oc-Bench benchmark would be a very demanding ablation given our limited hardware resources, and hence have left it for future work.
In the meantime, provide the results mentioned below. Table shows task-level performance is mostly consistent across ablations, although the "QMSum" task does seem to show some sensitivity to the number of retrieved tokens. This is further confirmed in Figure 14] which shows that longer examples benefit from more retrieved tokens. However, this is not the case in the other tasks, which seem to instead prefer less retrieved tokens. Furthermore, in these tasks, examples get shorter the best performing number of retrieved tokens also seems to decrease. This is consistent with our observations concerning diluted attention and the decrease in accuracy when attending tokens. Overall, such results, along with our positive results on the much longer oo- so-Bench benchmark, further confirm that our approach generally capable of efficiently handling context lengths much larger than the number of tokens available to the model at any one time.
many
Task
1K
2K
4K
6K
GovReport
31.26
31.44
31.33
31.26
QMSum
23.24
23.68
24.47
24.30
MultiNews
26.63
26.67
26.59
26.61
SAMSum
42.48
43.38
42.67
43.01
Table 12: Ablation of the size of the retrieved buffer (number of tokens) for EM-LLM; on Mistral with 4K local tokens on LongBench's 's summarization tasks.
34

---

Published as a conference paper at ICLR 2025
E FURTHER DISCUSSION POINTS
E. DETAILED ANALYSIS OF EM-LLM'S CONNECTION TO HUMAN EPISODIC MEMORY
This section provides a detailed analysis of how EM-LLM relates to human episodic memory (EM), addressing both the fundamental similarities and current differences between our computational approach and biological episodic memory systems.
1. Foundation: Transformers and Episodic Memory Integration
Recent work has shown that Transformer architectures naturally exhibit capabilities that parallel aspects of human episodic memory. In their basic operation. Transformers combine multiple pieces of information into coherent representations through latent embeddings recollections of concepts that have been inferred from inputs and encoded in the embedding space. These concepts can be recalled and utilized via the SoftMax self-attention mechanism, enabling human-like behavior in short-context recall tasks (Ji-An et al. 2024).
However, two fundamental constraints prevent Transformers from maintaining this connection to human episodic memory long-context scenarios: (a) The quadratic increase in computational and memory complexity and (b) the degradation of retrieval performance due to attention dilution Tworkowski et al. (2023); Ye et al. (224).
2. EM-LLM's Approach to Human-lite like Memory Processing
Our architecture extends Transformers' inherent memory capabilities beyond these limitations through two key mechanisms that mirror human cognitive processes:
2.A. Event-based Memory Formation We employ Bayesian surprise for event segmentation, a choice deeply grounded in both behavioural and neuroscientific evidence. Studies have shown that surprise signals in the hippocampus and other brain regions are crucial for event boundary and episodic memory formation (Sinclair et al., 2021 Zacks et al. 2007 2011
Sherman et al. 2022
Mariola al. 2022 Fountas al.) 2022). Our implementation demonstrates:
1. Content-Dependent Parsing: Analysis shows that tokens within surprise-based segments
exhibit significantly higher key similarity than tokens across segments (Table 2). indicating natural semantic grouping that aligns with human event perception.
2. Integration of Episode Components:
First, our model preserves temporal and contextual information within events, such as when, where, what, how, and who, as evidenced by strong performance on QnA and retricval tasks. In addition, while our current implementation focuses on text, the architecture is fundamentally compatible with multi-modal processing. Like recent multi-modal models, including Qwen2-VI (Wang et al. 2024), Pixtral (Agrawal et al.] 2024). LLaMa 3.2, EM-LLM can integrate different modality encoders into single embedding space, treating all embeddings equally in the KV cache.
2.B. Human-like Information Retrieval Our two-stage retrieval process combines both similarity- based retrieval for cued recall and decaying temporal contiguity reflecting free recall patterns. This integration enables our model exhibit both temporal contiguity effects and temporal asymmetry in recall behavioural patterns consistently observed in human EM retrieval studies (Howard and Kahana 2002). The inclusion of the contiguity buffer specifically allows for the maintenance of temporal relationships in way that mirrors human memory access patterns. Moreover, retrieval is done individually per-layer further supporting the transformer's 's learned ability to focus on different aspects of the sequence (see Appendix D.3),, including necessary contextual information. The combination of contiguity and layer-wist wise retrieval results a sophisticated information retrieval process which, combined with transformer, has all the tools allow for the complete contextual recollection of relevant events.
Current Limitations and Future Directions
While EM-LLM successfully implements key aspects of human episodic memory, several important differences remain:
1. Non-parametre nature: Unlike human memory, which involves complex synaptic weight changes, our method relies on non norammrmit storage of key-value pairs.
2. Hierarchical event structure: Current implementation lacks the sophisticated nested event representations observed in human cognition (Baldassano et al., 2017).
35

---

Published as a conference paper at ICLR 2025
3. Cross-Modal integration: While architecturally compatible, current implementation doesn't 't fully capture the rich multi-modal integration characteristic of human episodic memories.
4. Memory Consolidation: The model lacks mechanisms for long-ten terg memory formation processes and systems consolidation observed in biological memory systems.
These limitations represent opportunities for future work rather than fundamental flaws in our approach. A parametric version of EM-LLM could help reduce LLMs' memory complexity to bare minimum by replacing the vector database and KV cache storage requirements with a neural approach (e.g., using the model in Ramsauer et al. (021). Additionally, developing hierarchical event representations and integrating memory consolidation mechanisms could facilitate continual learning and further bridge the gap between artificial and biological episodic memory systems. In conclusion, we believe to have provided strong arguments that our approach is capable of complete recollections of experiences, resembling EM in humans. However, we acknowledge that lack explicit empirical evidence that this is how the resulting model makes use of the architecture in order to achieve the results presented, and hence clarify that we only claim an EM-inspired approach, rather than an actual human-like EM process.
Relationship to Different Memory Systems
While our work draws primary inspiration from episodic memory systems, the relationship between different types of memory (episodic, semantic, working memory, and other systems) is complex and often overlapping. both in biological and artificial systems. Our approach focuses specifically on episodic memory-like features such as event segmentation and temporal organization of experiences. However, as discussed in Section [ our architecture also shows interesting parallels to working memory models, particularly in how the local context aligns with concepts like Baddeley's working memory system and Cowan's 's focus of attention. While some aspects of our model, particularly the learned representations in the underlying LLM, may share characteristics with semantic memory systems, the primary innovations in EM-LLM centre on episodic-like features. Future work could explore these connections more explicitly, potentially leading architectures that better capture the interactions between different memory systems, including the development of modality-specine buffers inspired by Baddeley's 's multi-component model.
F.2 ARCHITECTURE CONTRIBUTIONS OF EM-LLM
EM-LLM introduces three novel architectural contributions for LLMs, for which we have shown their importance both conceptually and with empirical results.
(1) Dynamic surprise-basse baaiee segmentation. The method to segment a given context window based on Equation (I) is the first method for dynamic segmentation of KV cache into blocks, and also the first that manipulates the KV cache based on insights from cognitive science, using an intrinsic measure to LLMs. We show empirically using multiple LLMs that this low-cost and method is able group relevant pairs of keys and values (KV) together (relevance measured as key similarity) with much higher accuracy than fixed segmentation, the only alternative proposed approach (See Table for key similarity comparisons). We also show that this method results in increased LLM performance, especially in retrieval tasks (16.6% average increase over InfLLM) and molti-docummnt QA tasks (6.4% average increase over InfLLM) across all the LLMs we tried (See the "S" column in the tables of Appendix [A..)
(2) Graph-based refinement. The method presented in Algorithm i[is the first to refine the temporal borders of events in the context window of LLMs using graph theory. We relied on the insight that tokens are more useful to be recalled together, if the variance between their keys is low, as they need to be used by a single query at the time. This method can also stand by itself dynamic segmentation approach of KV cache, more computationally heavy than surprise- based segmentation but achieving a competitive accuracy in grouping relevant (KV) together (see again Table 2), while it has the extra benefit that can be used in each attention head independently. without relying on the LLM output.
(3) Contiguity buffer. This is dedicated decaying buffer in the context window of LLMs that maintains the KV cache of temporally contiguous tokens to the context window of the LLM for a certain amount of time. This relies the recent insight that self-atte attentitoio heads responsible for in-context learning are shown to consecutively attend to contiguous groups, similarly to human studies (Ji-An et al., 2024). We show that this algorithm also be combined with methods (1)
36

---

Published as a conference paper at ICLR 2025
and (2) and results in further increases in the overall LLM performance. Notably, the average increase in retrieval tasks over InfLLM jumps to 19.4%, and for multi-document QA tasks to 9.2% across all the LLMs we tried (See the "SM+C" column in the tables of Appendix A.1
E.3 WHY IS THE REFINEMENT ALGORITHM EFFECTIVE?
The use of the argmax in Algorithm| I guarantees either a positive improvement or no change in similarity for each event boundary position update, hence either improving overall similarity or showing no change from surprise-based segmentation. Therefore, while we would ideally find the globally optimal event boundaries with regards to the similarity metrie, and seek to converge this point, this would be much more expensive to compute and introduce a lot of overhead for every processed chunk of the context and the corresponding memory units. Instead, our algorithm simply implements a cost-effective way to look for any potential increase to this metric, as it has been empirically shown to do successfully in section 4.2 Nevertheless, to briefly touch on the convergence of such a method, our approach can be seen as single pass of Phase 1 of the heuristic Louvain method (Blondel et al., 2008) initialized with surprise-based segmentation (as opposed to having each node assigned its own community), and modified to only consider the move of a node to its right-side neighboring community. As our initial results had shown that surprise- based segmentation already achieves higher similarity metrics (including modularity, which is the objective used in the Louvain method) than fixed or random segmentation (Table 2. we believe this is a good initialization as it means that our algorithm will, at worst, achieve the same similarity metries. While the Louvain method is considered to be an efficient way to converge to local optima when iterated, our own modifications and lack of iterations mean we cannot claim such behavior but rather suggest that we are likely to see some improvements in our metrics, as our results have confirmed.
E.4 FEASIBILITY OF TEND-O-EEN NEURAL IMPLEMENTATIONS
A significant difference between EM-LLM and biological memory systems is the ability of neural circuits in the brain to learn and adapt their event segmentation and memory retrieval mechanisms through experience. In artificial neural networks, this would correspond to end-to-end optimization via differentiable architectural components. Below, we discuss the feasibility of such an approach and compare it with our current implementation:
Event segmentation: Differentiable event segmentation models have already demonstrated the feasibility of learning a temporal structure from continuous experience. Models like SEM (Franklin) et al. 2020) show how neural networks can combine with probabilistic inference to capture human- like event segmentation, while approaches like the DLH (Zakharov et al. 2022b) demonstrate that neural architectures can learn to identify hierarchical temporal boundaries through differentiable clustering and amortised variational inference. For instance, using the VaDE trick in (Jiang et al., 2016). These approaches offer powerful advantages in terms of learned representations and flexibility, potentially capturing the complex hierarchical event structure of real environments and adapting to different domains. Particularly compelling advantages include the ability perform layer-wise or attention-head hww-hhhhaww--e wise segmentation and the potential emergence of nested timescale structures, as demonstrated in (Zakharov et al., 2022a bb). mirroring how the brain processes events at multiple temporal scales (Baldassano et al. 2017). While such end-to-end training is theoretically appealing and mirrors how neural circuits might learn temporal structure, our method takes more pragmatic approach by leveraging the pre-trained capabilities of LLMs. By using Bayesian surprise computed directly from model outputs to detect event boundaries, we achieve efficient segmentation without requiring complex architectural modifications or additional training, while still aligning with cognitive theories about prediction errors in event perception (Zacks et al. 2007).
The development of neural architectures for memory retrieval has evolved from classical Hopfield networks (Hopficld) 1982) through several key innovations. Early Hopfield networks demonstrated how content-addressable memory could emerge from simple neural circuits, paralleling biological memory systems. This was significantly advanced by Neural Turing Machines (Graves) 2014) and their successor, the Differentiable Neural Computer (Graves et al., 2016),, which introduced differentiable memory access mechanisms. Modern Hopfield networks (Ramsauer al.) 2020) further revolutionized our understanding by establishing a theoretical connection between transformer attention and associative memory, showing how these systems store and retrieve exponentially many patterns while maintaining stable dynamics. Such benefit the quality of memory representations, they could learn optimal projections for generating
Retrieval:
approaches could particularly
37

---

Published as a conference paper at ICLR 2025
representative keys for memory blocks, potentially capturing universal contextual patterns more effectively than our current approach. While end-to-end training of memory systems is feasible, as demonstrated by models like MERLIN (Wayne et al. (018), such approaches often face challenges with credit assignment over long sequences and require complex architectural modifications. Our KNN-based approach leveraging the KV cache offers pragmatic middle ground: it hamnesses the rich semantic representations already present in transformer models while maintaining the computational benefits of nearest- -niihbor neighbour retrieval. This aligns with both biological intuitions about pattern matching in the hippocampus [O'Reilly and Norman 2002) and the theoretical foundations of modern Hopfield networks, where similarity-basee attention serves as a form of associative memory. By operating pre-trained representations, our method sidesteps the training complexities of fully differentiable memory while preserving the benefits of content-based retrieval.
The refinement of event boundaries could also theoretically be learned end-to-end, similar to how attention pruning mechanisms (Ying et al. 2019) learn to identify optimal subgraphs in graph neural networks, or how hierarchical clustering can be made differentiable (Ying et al. (2088). Our graph modularity approach provides computationally efficient alternative that optimizes for coherence within segments while respecting the initial surprise-based boundaries. While our method primarily motivated by computational considerations, parallels how memory consolidation might strengthen associations between related elements within event while weakening cross- event associations (Preston and Eichenbaum 2013). The high modularity of our surprise-bar based segmentation, even before refinement, suggests that prediction errors naturally tend to occur boundaries between coherent event structures.
Refinement:
E.5 FUTURE EXTENSIONS INSPIRED BY HUMAN MEMORY SYSTEMS
Human episodic memory exhibits several sophisticated features beyond those currently implemented in EM-LLM. Here, we discuss how incorporating these additional characteristics could enhance our model's 's capabilities and performance:
Hierarchical organisation: A hierarchical structure in memory can provide multiple advantages such as improved retrieval, more disentangled latent embeddings, longer future predictions (Saxena al.) 202 Zakharov et al. 2022b), better planning (Hafner al.) (022) and higher agreement with neural processes the brain Baldassano et al. (2017). our model, a hierarchical organisation of episodic memories based on the existing hierarchy of embeddings in the LLM layers could be implemented by extending our segmentation processes to operate at each layer of the Transformer independently. This could be achieved either through a differentiable approach a layer-specific surprise metric. Interestingly, our current k-NN retrieval approach already implicitly leverages hierarchical structure through its underlying approximate nearest neighbour algorithms, which typically employ tree-based structures (Johnson et al. 2019) to efficiently partition the embedding space.
Memory consolidation: The brain's process for memory consolidation is crucial for continual learning, ability that remains largely unsolved in current LLMs. Implementing consolidation mechanisms in EM-LLM could help address catastrophic forgetting while enabling more efficient integration of new information with existing knowledge.
Mental time travel: The ability to employ the same retrieval mechanism for imagining future events for recalling past experiences is key feature of episodic memory that could significantly enhance LLMs' planning and reasoning capabilities. By leveraging its event-based structure to simulate potential future scenarios or recall past experiences in novel contexts, this mechanism could provide powerful solution for planning and reasoning, which are currently important challenges in large generative models.
F PROOFS
F.1
APPROXIMATE EQUIVALENCE OF K-NEAREST NEIGHBOURS AND SOFTMAX ATTENTION
Here will attempt to show that using k-NN retrieval in key-value cache part of the attention mechanism in transformers is an approximation of applying softmax attention over the entire sequence of tokens.
38

---

Published as a conference paper at ICLR 2025
Let q be query vector and K {k1, ......, ) the set of key vectors in transformer model with dimensionality Each key k; has a corresponding value vector 8s with V {11 U2,..., Un The softmax attention weights a, are defined as:
exply ki d=
04 Σラ= expla Ri d=
(6)
The output vector u is computed as:
u
(7)
In the k-NN approach, subset K' of size k is selected, containing keys nearest The approximated attention weights of over this subset are:
expla ki dd
df EjEK' exp(y k2 d-
for kg E K'
(8)
The approximate output vector u' is:
= Σ K,EK'
ai
(9)
ASSUMPTIONS
1. Exponential Dominance: The exponential function in the softmax is sharply peaked, imply- ing that keys with the highest similarities to contribute significantly more to the sum than others.
2. Representativenes: of k-NN Subser: The subset K' captures the majority of the attention weight from full set K.
Lemma 1: Dominance of k-NN Subset If K' consists the ls keys with the highest dot products q ki, then:
Σ
expla ki d 2) a explq k3 2)
(10)
JEK'
for some a As 1, typically very close to 1.
Proof: This follows from the exponential dominance assumption and the nature of the exponential function, which is sharply peaked.
Approximation of Output Vector Given the dominance of K' as shown in Lemma 1, Lemma 2:
the approximate output u' effectively represents the full output u:
||
ull <c
(11)
where e is a small error term.
Proof: Follows from the weighted sum structure of u and u', using the bounds established in Lemma 1.
Given the lemmas and under the stated assumptions, the k-NN retrieval mechanism within key- value cache effectively approximates the softmax attention mechanism transformers. This proof highlights the efficiency versus accuracy trade-off inherent in using approximate methods like k-NN retrieval.
39


Reviews:
[{"review_id": "RgNDruhXts", "reviewer": "Reviewer_SyTr", "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.", "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. \n2. To better justify the claim that the mechanism is EM-related with further analyses.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "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.\n\nAnother 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gbE9tWRzbO", "reviewer": "Reviewer_pMUj", "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.", "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.\n- 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? \n- Clarifications for Figures:\nIn Figure 1, does the data reflect results from S, SM, or SM+C? \nIn Figure 2, the memory units seem to have variable sizes, different from the fixed size in InfLLM? Additionally, how are the bars normalized?\nIn Figure 4A, S appears closer to human segmentation than SM or SC, contrasting with Figure 4B. \nIn Figure 4, what does SC stand for? Does \"C\" refer to the continuity buffer or represent something else?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "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.\nQuality: 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.\nClarity: 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.\nSignificance: 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.  \n- Although InfLLM is a primary benchmark comparison, it is absent from Figure 1. Including InfLLM in Figure 1 would provide a more complete comparison.\n- 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.\n- In Formula (1), notations (\\mu_{t-\\tau:t}) should be corrected, keeping consistent across the context.\n- Certain figures would benefit from clarification, as detailed in questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "w33ZsdxOzE", "reviewer": "Reviewer_RkzP", "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.", "questions": "1.How sensitive is the method to the choice of the surprise threshold parameter? What is the process for tuning this parameter?  \n2.How do you ensure the stability of the boundary refinement process? Are there cases where it fails to converge?  \n3.Does the issue of needing to query too many matching events still exist when the context window is too long?  \n4.If we use the Surprise-based method to segment the context into a RAG format, will the performance remain the same?  \n5.How does the method handle dynamic changes in context relevance over time?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "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;  \nThe use of Bayesian surprise for dynamic event segmentation in LLMs is novel, moving beyond simple fixed-length segmentation used in prior work like InfLLM;  \nThe method achieves 100% accuracy on Passkey.Retrieval task with sequences up to 5M tokens, demonstrating practical scalability;  \nThe 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.  \n2.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.  \n3.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.  \n4.Although Table 1 shows some improvements over InfLLM, the improvements on most tasks may not be statistically significant.  \n5.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.  \n6.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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4l8RMvkZMc", "reviewer": "Reviewer_N12r", "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", "questions": "The same to Weaknesses.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "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?\n\n2. 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.\n\n3. 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?\n\n4. In Eq. (2), could you explain how you derive the key vectors?\n\n5. You employ k-NN in the MEMORY RETRIEVAL stage—does this significantly increase computation time?\n\n6. I found Figure 4 difficult to understand. Could you provide additional explanation to clarify it?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)