AI Agent Memory: 6 Ways Systems Retrieve the Right Context

Quick Summary

This article works through the main ways an AI memory system decides which stored memories to put into an LLM’s context: keeping everything in the context at all times, a memory bank with an LLM-generated semantic index, embedding-based semantic search, classical word matching, graph traversal, and scoring methods that combine several of these. It sets out the bottleneck behind each approach and what determines which one fits a given system.

Why AI systems need memory

As personalized AI systems are becoming more and more involved in our lives, the need for AI agent memory is becoming apparent. When it comes to an AI assistant or tool adapting to our needs over time, memory is the key. To better understand this, let's first tak a look at how AI systems work roughly.

In almost all applications, AI technology depends on what we call Large Language Models (LLMs), which are statistical representations of human laguage. This representation is derived from very large datasets (like the contents of the World Wide Web) and is stored as weights determining the properties of the LLM. The tuning of these weights is conducted by a process of Machine Learning, which gradually fits the parameters of the model through training on the large datasets.

This training is often carried out by companies with lots of resources. However, these so-called foundational models which are the results of training have no information stored about you, the user, only what is available in large public datasets. This is where memory comes into play.

Memory and the LLM context

What an LLM does is taking a string of information chunks called tokens as input, and returns another string of tokens as output. The input based on which output tokens get generated is called the LLM's context. Any queries sent to the AI application and any additional instructions the AI system is given is passed into the context.

Information about the user is limited compared to the information required to model human language, so the LLM's context is a natural place to store the former. So to instate a memory system the question to answer is when and how the right information gets passed to the context. This is basically the fundamental question in so-called Retrieval-Augmented Generation (RAG) systems.

Figure 1: Everything the model knows about the user arrives through the context, one query at a time.

6 Approaches to memory selection

1. Keeping everything in the context

The simplest approach to memory selection is simply to store every bit of information in the context all the time, with every query. This way, there is no need to devise a clever way to match the user queries to the right piece of information, the LLM takes care of this implicitly with every LLM forward pass - it reads every stored memory and decides through its inherent reasoning capabilities, which piece of information is relevant and needed to produce a correct answer.

If you don't expect to have a lot of memories, this implementation might be just the right fit for your system.

How much is too much and why?

This is something for which there is no well-grounded theoretical answer and is something that you can derive empirically for a given system. Generally speaking, the more unrelated memories there are in the context, the harder it is for the LLM to find the right entries and the more noise there is, the less computational power is left for completing the task at hand.

2. Memory banks and a semantic index

Because context token capacity is a limited resource it can make sense to implement a different strategy for memory management and have a separate system of memory selection. In this case, memories are stored in a central memory bank and there is a mechanism for selecting which pieces of information to add to the LLM's context for each query.

One way to do this is by reducing the amount of context the memory bank requires through a more compact representation of the memories. This compression can be achieved by creating summaries, or a semantic index if you will. The LLM context gets the index and decides which memories to add in full into the final LLM pass. A system that implements this approach successfully is Claude Code. In this solution, the selection of memories is still up to the LLM's cognitive capabilities.

The bottlenecks in this case are the time it takes for the LLM to produce output tokens (forward pass), and the time it takes for maintenance of the semantic index. In addition, the quality of the index depends on the ability of the LLM to produce meaningful summaries of larger memory items. Since the LLM is a black box, it becomes a challenge to have guarantees for the quality of outputs it produces.

Figure 2: The index is read first, then the chosen memories are fetched in full for the final pass.

3. Embeddings and semantic search

Another way to select memories for LLM context is to leave the LLM out of the picture altogether.

How then, can we select the right text for the right query?

A possibility is to open this black box and use some knowledge of the internal mechanism of the LLM for our use. As an input text passes through the LLM, it gets transformed into a numerical representation, a vector with a lot of dimensions. In other words, the text gets embedded into a vector space.

A vector space is a mathematical concept with the property that it makes sense to talk about the distance of its elements. In case of the LLM's vector space, the crucial property is that to some extent, texts that we consider closer in meaning are represented by vectors that are closer in the vector space. This way, semantic similarity is translated to a mathematical representation. We can use this representation to perform searches within a memory bank that takes the meaning of the texts into account.

A tool that transforms texts into their vector representations is called an embedder. A so-called semantic search can be performed in a memory bank that consists of the embedded vectors of the memory texts. An advantage of this solution is that a database of number arrays can be indexed with sophitisticated methods for fast retrieval. The bottleneck here is the ability of the embedder to succesfully represent semantic similarity with embedding similarity.

Figure 3: Semantic similarity becomes distance. Shown in two dimensions; real embeddings have hundreds.

4. Word matching algorithms

Other methods for searching similarity between texts are more classical word matching algorithms. Although these are combinatorially more expensive, they can be simplified and combined with embedding based methods.

5. Graph-based memory navigation

A further way to navigate among memories is by introducing another mathematical concept to represent the memory bank: a graph. Graphs represent connections between objects, resulting in a network of nodes, in this case, memory texts.

If two memories are semantically connected in some way, this can be represented by an edge in the graph. Travelling along the edges of a graph gives a new meaning of distance between its elements, in a way representing semantic distance.

Figure 4: Edges record semantic connections, so distance can be counted in hops rather than similarity.

6. Combining approaches through scoring

These different approaches can be combined through scoring the elements of the memory bank and returning the elements with highest scores. Hindsight is an example of a system that combines the differrent distance measures through weighting.

The approaches side by side

Choosing an approach

To sum up, we have systematically explored some of the ways by which memory systems can select the right memories for the right user request. Which solution is the right pick for you depends on performance requirements of your system and your confidence in the different language modelling technologies currently available. Which one will work for you is difficult to answer, it is up to you to test them out.