Human-like Episodic Memory for Infinite Context LLMs

Recently, several approaches have been proposed to extend the context window of Transformer-based models. Those include methods that address the limited representational capacity of softmax attention, and its quadratic computational and memory cost [Katharopoulos et al., 2020, Munkhdalai et al., 2024]. Other methods target the poor extrapolation of typical positional encodings to out-of-distribution (OOD) 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 popular Rotary Positional Embeddings (RoPE) [Su et al., 2024]. To address this, some approaches propose scaling of the rotation angles [Chen et al., 2024a] or the base constant [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], exploring alternative strategies such as KERPLE [Chi et al., 2022] and FIRE [Li et al., 2024] or adopt mechanisms from certain LMs like T5 [Raffel et al., 2020].

Concerning computational efficiency and diluted attention, successful approaches propose methods for general improvements to the efficiency of Transformers through optimised computations [Dao, 2024, Han et al., 2024a, Aminabadi et al., 2022, Kwon et al., 2023, Liu et al., 2024c, Brandon et al., 2023] or compression techniques [Nawrot et al., 2024, Zhang et al., 2023], as well as training methods tailored for long-context scenarios [Zhu et al., 2024, Chen et al., 2024b]. Another direction is the utilisation of retrieval-based methods, the vast majority of which relies on a vector database that keeps a 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 a key-value cache with k-NN lookup can be seen as an approximation of applying softmax attention to the full token sequence (see Appendix A.4), k-NN retrieval methods can be used without any fine-tuning [Bertsch et al., 2023]. For an exception that does not rely on k-NNs, see Wang et al. [2023].

A 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 InfLLM [Xiao et al., 2024a], which achieves state-of-the-art performance on long-context benchmarks. InfLLM 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. This can be seen as a form of hierarchical attention, as illustrated in Fig. 1. While group-based retrieval represents a promising direction, our approach significantly advances this concept by dynamically determining token groupings in a manner akin to human memory formation, addressing a fundamental limitation of InfLLM’s fixed-size segmentation and enabling more adaptive and context-sensitive processing of extended information.

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 experiences and suggesting links between memory, efficient representations and navigation of physical and conceptual spaces [Gershman et al., 2012, Benna and Fusi, 2021]. In machine learning, episodic memory-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., 2016, Pritzel et al., 2017, Coda-Forno et al., 2024]. In addition, models of memory construction and consolidation have been successful in alleviating catastrophic forgetting in neural networks [Kirkpatrick et al., 2017, Lopez-Paz and Ranzato, 2017, Chaudhry et al., 2019, Buzzega et al., 2020, Prabhu et al., 2020], including LLMs [Das et al., 2024], and appear to 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 experiences are segmented into events, and when new memories are encoded and retrieved [Lu et al., 2022]. Surprise plays a critical role in this process, triggering event boundaries and memory formation [Fountas et al., 2022, Kumar et al., 2023]. 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., 2021] 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.2A). 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.2B) 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-based event segmentation and LLMs’ inherent temporal contiguity and asymmetry effects to enable a new generation of Infinite Context-Length LLMs, capable of processing and understanding information over vastly extended timescales.

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 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, 2001]. The evicted tokens typically comprise the majority of past tokens in a long-context scenario, extending far beyond the LLM’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, as first observed by Xiao et al. [2024b], Han et al. [2024b] and later adopted by Xiao et al. [2024a]. For retrieved tokens, which are therefore discontinuous and outside the local context, we assign a fixed position embedding as in Raffel et al. [2020], Xiao et al. [2024a]. This architecture enables EM-LLM to effectively process and utilise information from positions outside its pre-trained local context window, while maintaining the underlying LLM’s performance characteristics.

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 as an 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 an auto-regressive model, with high values indicating the unpredictability or novelty of each new token within the context according to the model, i.e., it is ‘‘surprised’’ by the next token. Following work on cognitive modelling [Roseboom et al., 2019, Fountas et al., 2022], we employ a thresholding mechanism to perform an initial identification of event boundaries (used for the first time in LLMs). Formally, a token $x_{t}$ is considered a potential boundary if its surprise value exceeds a threshold $T$:

where $\mu_{t-\tau:t}$ and $\sigma_{t-\tau:t}^{2}$ are the mean and variance of surprise for a window offset $\tau$, and $\gamma$ is a 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 $\gamma$. We also explored a fixed threshold approach ($T=T_{\text{fixed}}$), though our primary focus remained on the dynamic threshold due to its adaptability to varying contexts. This initial segmentation results in a set of potential event boundaries $\mathcal{B}={b_{1},b_{2},...,b_{k}}$, where each $b_{i}$ represents the index of a token exceeding the surprise threshold. These boundaries serve as the starting point for our subsequent refinement process, which aims to optimise the intra-event coherence and inter-event distinctiveness of the resulting memory segments.

While surprise-based segmentation provides an effective initial estimate of event boundaries, we make the key observation that the utility of elements within an event during memory recall depends 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. 1. To further ensure this, we introduce a boundary refinement step which looks to optimise this objective. Such an objective is typically optimised in the context of graph-clustering, hence we will express this refinement process in a graph-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 $x_{1},x_{2},...,x_{n}$ as an adjacency matrix. We define the adjacency matrix $A^{h}$ as

where $K^{h}_{i}$ and $K^{h}_{j}$ are the key vectors corresponding to tokens $x_{i}$ and $x_{j}$, respectively. The similarity function measures the closeness of two key vectors; in our implementation, we use dot product similarity ${K^{h}}^{\mathsf{T}}_{i}\cdot K^{h}_{j}$ 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 $f(A,\mathcal{B}):\mathbb{R}^{n\times n}\times\{1,\ldots,n\}^{k}\rightarrow\mathbb{R}$. 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 $\mathcal{B}$. We experimented with two widely-accepted graph-clustering metrics: modularity and conductance [Miasnikof 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 a random cluster. As our edge weights represent the similarity between two tokens, we seek to maximise this metric. Modularity is defined as:

where $m$ is the total edge weight in the graph, $c_{i}$ is the community (episodic event) to which node $i$ is assigned, and $\delta$ is the Kronecker delta function. Conductance, on the other hand, measures the fraction of total weighted edges cut by a given community boundary, and is defined as:

where $S=\{b_{i},b_{i}+1,...,b_{i+1}\}$ is a subset of all nodes $V=\{b_{1},b_{1}+1,...,b_{k}\}$ in the induced graph, with $b_{i}\in\mathcal{B}$. Lower conductance values indicate better community structure. Our boundary refinement algorithm iteratively adjusts the initial surprise-based boundaries to optimise these metric functions. While our best results were achieved using modularity, we also include comparisons with conductance-based boundary refinement to provide a comprehensive analysis. The overall process can be summarized in Algorithm 1.

This algorithm first identifies initial boundaries based on the surprise threshold $T$, then refines these boundaries by finding the optimal position $\hat{\beta}$ between each pair of consecutive initial boundaries $(\alpha,\beta)$ 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 1 is dominated by the boundary refinement step, which has an overall complexity of $\mathcal{O}(kn)$, where $k$ is the number of initial boundaries and $n$ is the sequence length. A detailed analysis of the algorithm’s complexity, including the computation of modularity and conductance metrics, is provided in Appendix A.2. Despite this modest computational overhead, the resulting improvement in segment quality leads to significant performance gains in downstream tasks, particularly those requiring complex temporal reasoning.

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 a two-stage mechanism to select relevant episodic events for the LLM’s context window (Fig.2C). First, we retrieve $k_{s}$ events using $k$-NN search based on 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_{s}$ events, retrieved based on their similarity to 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 $k_{c}$, this buffer promotes temporal relationships in retrieval. When an event is retrieved, we also enqueue its neighboring events (within $\pm n$ 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 a natural decay of temporal context as new events are processed, with older or repeated events being dequeued as new ones are added. In total, $k=k_{s}+k_{c}+2$ events are added to the context window, striking a balance between relevance and temporal relationships in a manner analogous to human episodic memory retrieval.

As previously mentioned, InfLLM is, at the time of writing, considered to achieve state-of-the-art performance on long-context benchmarks ($\infty$-Bench, LongBench), as well as being the only method which uses group-based k-NN retrieval in LLMs on such benchmarks. We, therefore, employ this model as our baseline for comparison with our own methods on short context windows (4K+2K, as in Xiao et al., 2024a).

Results on the LongBench dataset (Table 1) show that our method is able to improve on InfLLM in all but one task, with the best method achieving an overall increase in performance of $1.8$ percentage points (a relative improvement of $4.3\%$). Note that the table shows the best single method in terms of overall performance for each ablation. Looking at individual task performance across all experiments, we are able to beat InfLLM in all tasks (see Appendix A.1). Interestingly, we see an especially large jump in performance on the PassageRetrieval task across all ablations, with up to a $33\%$ improvement on InfLLM. This task requires the model to identify the original paragraph from a summary, a challenging task that tests the model’s ability to accurately recall a wide range of detailed information from a large context concurrently. The substantial improvement on this task highlights the effectiveness of our event segmentation method in enhancing long-term memory recall and retrieval accuracy in LLMs. Additionally, our method achieves a notable $9.38\%$ improvement on the HotpotQA task, which involves complex reasoning over multiple supporting documents, further emphasising the model’s ability to provide explanations for answers.

As mentioned in Section 3.2, we employ modularity and conductance as two refinement objectives in our boundary refinement algorithm, due to 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:

Kumar et al. [2023] found strong correlations between human-perceived events and prediction errors across 3 short podcasts (7-30 minutes on average), when processing the corresponding transcript with an LLM. Taking advantage of such human data and results from previous works on this dataset [Michelmann et al., 2021, Lositsky et al., 2016], we compare the segmentation quality and correlation with human data for each of our methods (3) using our similarity metrics.

As shown in Fig. 3A, human-perceived 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 ($\mathbf{S}$) achieves very similar results to humans, while the addition of our refinement algorithm ($\mathbf{SM}$, $\mathbf{SC}$, $\mathbf{FM}$, $\mathbf{FC}$) significantly improves performance. Fig. 3B further shows that surprise-based methods ($\mathbf{S}$, $\mathbf{SM}$, $\mathbf{SC}$), consistently identify event boundaries that are closest to those perceived by humans.

Looking at Table 2, it is clear that surprise-based segmentation with refinement ($\mathbf{SM}$, $\mathbf{SC}$) provides the best results in terms of event similarity metrics, regardless of the base LLM used. While the surprise-only method ($\mathbf{S}$) achieves some good results, we observe that refinement 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 ($\mathbf{FM}$, $\mathbf{FC}$) do not reach the same performance as their surprise-based counterparts, further showing that the initial segmentation with a surprise threshold is crucial to achieving the best possible balance in intra-/inter-similarity with our methods.

As demonstrated in Tables 1 and 2, along with Fig. 3, each of our ablations show various positive improvements on InfLLM (also see Appendix A.1). As mentioned in Section 4.3, refinement has a strong positive impact in improving our similarity metrics. This is seen to translate well to model performance in Table 1, achieving the best performance in a third of the tasks, as well as agreeing with human data (Fig. 3). The effects of contiguity are also clearly demonstrated in this table, with the addition of our contiguity buffer achieving the best performance on three tasks, and the second-best overall score. Furthermore, these methods are shown to be generally complementary, achieving the best overall performance when combined.

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 to the nature of the tasks and the varying importance of contiguity across these tasks. For instance, in Supplementary Fig. 5, the MultiNews task scores higher than our baseline only for a ratio of $70\%$ contiguity to similarity buffers. 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 a contiguity buffer that is as big or smaller than the similarity buffer yields the best results, 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.