← Posts

Memory in AI Agents

Solving memory is arguably the most important problem in getting AI agents to work as well as humans.

By 'solving memory', specifically, I mean developing an efficient and performant way to embue AI agents with cross-session memory that is persistent.

Today, most AI agents don't persist learned attributes or features across user sessions in a way that makes the human-to-agent interaction feel natural.

LLMs such as ChatGPT and Claude can remember certain things about you such as your name, some of your preferences and other smaller things. But it can't remember that you asked it about X 5 days ago and then recall that conversation. At least nowhere as well as a human can.

This might not be a big problem if we had unlimited context windows. One thought is that we would continue to add things to the context window and maybe perform some pruning over time. But this has a few problems. First, this would lengthen inference times and impact the user experience. Remember that every token in the input must be embedded, normalized, attended to, normalized again and so on. As of yet, we don't have an O(1) way of running this process, so naturally the inference time is going to increase as the number of tokens increases.

Second, there are a lot of details in a conversation that just aren't important and we likely wouldn't want to carry forward anyways. While our brains aren't perfect, they're quite good at identifying what is important to remember and what isn't.

Intuitively I would describe human memory as a hierarchical strtucture that progressively gets noisier as you go deeper into the structure, and that noise causes you to misremember details or memories altogether. There has been a lot of interesting work done around how traumatic experiences shape memories and how, for example, we have better recall for memories that were shaped in situations with a lot of adrenaline.

To reproduce something like this in an LLM, we would likely need to add some sort of "memory coefficient" for certain memories that, regardless of how far down the structure they are, can be remembered precisely.

Today, in LLMs, we can create crude "memories". One way that I've done this is with a tool call in ChatGPT that says "if the user proviodes an important fact about themselves, save it." Here is what it looks like in code:

		Tools: []openai.ChatCompletionToolParam{
				{
					Function: openai.FunctionDefinitionParam{
						Name:        "save_memory",
						Description: openai.String("Save personal information about the user that would be
                        helpful to remember in future conversations. This includes: location/address,
                        preferences, family details, important dates, interests, or any personal
                        facts the user shares. Use this whenever the user mentions something
                        personal about themselves."),
						Parameters: openai.FunctionParameters{
							"type": "object",
							"properties": map[string]interface{}{
								"memory": map[string]interface{}{
									"type":        "string",
									"description": "The personal information to remember
                                     about the user. Be specific and include context.",
								},
							},
							"required": []string{"memory"},
						},
					},
				},
			},

If I prompted the agent with "What is the weather in San Francisco?", it might create a memory and save it to a database like this:

memories

Over time there would be a list of memories and they would get loaded into the context window. Like I said, this is pretty crude. You can easily see all of the infrastructure that would need to go into supporting this. For example, compaction of memories, prioritization, and much more.

I think in an ideal world, the LLM would recursively update a memory module, similar to how the hippocampus and neocortex work to write, index and distribute memories across the cortex. If this is a separate module outside of the LLM, then we can probably evolve our existing crude setup to mirror this. For example, on every write to the memories table, we run a post-write trigger that compacts memories, re-organizes and re-indexes them. I wonder how far that would get us?

Alternatively, if it needs to be "part" of the LLM i.e. stored in some subsection of a hidden layer(s) somewhere then we would need to figure out how to recursively update those weights.

Clearly, still a lot of work to be done here but I think folks are waking up to the fact that memory performance (efficiency and accuracy) is increasingly the bottleneck to geting human like behavior out of LLMs.