More Than Reasoning: What Memory Experiments Taught Me About Language Models

How a project about managing an AI agent’s memory turned into an investigation of what language models actually do with the information we give them.

Why do agents suck sometimes?

That was pretty much the question I started with.

We have AI agents that can browse the internet, read documents, search through conversation histories, and even store information for later. Yet somehow, they still manage to fail at questions that seem ridiculously straightforward.

Naturally, I wanted to understand why.

One of the first things I became interested in was memory. More specifically, how an agent decides what information to remember, what to forget, and what to retrieve when it needs to answer a question.

I figured that if we could give an agent the right information at the right time, we could make it perform better.

Sounds reasonable, right?

Well, somewhere along the way, I ran into this rather interesting example.

WHAT!?

The answer was literally sitting right there! Heck, the model had even written it down in its own note before deciding that it didn’t know the answer.

For some context, this was Qwen2.5-3B-Instruct^[3], a relatively small language model with around three billion parameters. It was given eight selected memory entries from a conversation in the LongMemEval-S benchmark^[2].

We had instructed it to write a brief note before answering, and to respond with “unknown” if the memory didn’t contain enough information.

But here’s the interesting part: the memory did contain the answer, and the model had reproduced it.

Now, before we get too excited, this doesn’t necessarily mean the model understood everything and simply decided not to answer. There are several possible explanations for what happened.

What it does tell us, however, is that retrieving the correct information doesn’t guarantee that a model will actually use it to produce the correct answer.

That observation would eventually change the direction of my research.

But to understand how I got here, we have to go back to where this project started.

1. Can we treat memory as an optimization problem?

One of the classes I took at UC Berkeley was CS170, Efficient Algorithms and Intractable Problems.

A big part of the class involves taking seemingly complicated problems, figuring out their mathematical structure, and designing algorithms to solve them efficiently.

So naturally, I wondered whether we could do something similar with memory management.

Imagine you’re having a conversation with an AI assistant over several months.

You tell it about your classes, your favorite coffee, the restaurants you’ve visited, the projects you’re working on, and a whole bunch of other things.

Some of that information might become useful later. A lot of it probably won’t.

The problem is that language models have a limited amount of context they can process at once. We can’t necessarily keep feeding the entire conversation history back into the model every time we ask it a question.

We have to be selective.

In my initial framework, I divided memory into three categories:

1. Active memory: Information currently visible to the language model.
2. Archived memory: Information stored somewhere else that the agent could retrieve later.
3. Deleted memory: Information that’s permanently gone.

The goal was straightforward: preserve the information most likely to help answer future questions without exceeding our memory budget.

Of course, there’s a catch.

How do we know which information will be useful in the future?

We don’t.

But what if, just for an experiment, we pretended that we did?

Building an oracle that knows the future

This is where my amazingly intelligent friend Claude and I started working on a hindsight optimizer.

An oracle is essentially a system with privileged information. In our case, it knew exactly which pieces of memory would be required to answer future questions.

Think of it like studying for an exam, except someone has already handed you the questions.

You’d know precisely which topics to focus on and which ones you could safely ignore.

For our restricted keep-or-delete problem, we implemented an exact optimizer with the following objective:

Okay, that’s a mouthful. Let’s unpack it.

Here, $y_q$ indicates whether we preserved all the evidence needed to answer question $q$. Meanwhile, $x_{i,k}$ tells us whether memory item $i$ is still active at its $k$-th required use, and $s_i$ is the number of tokens that item occupies.

The first term rewards us for preserving enough information to answer future questions. The second term gives us a small incentive to use less memory.

We also introduced constraints to make sure our optimizer behaved realistically.

For instance:

This says that once we’ve deleted something, we can’t magically bring it back.

Another constraint ensured that every piece of evidence required by a question was covered by at least one retained memory item:

This applies to each evidence requirement $r$ of question $q$, where $C(r)$ represents the memory items containing that evidence.

Finally, we enforced our active-memory budget:

Here, $B$ is the maximum number of tokens we’re allowed to keep active, and $\tau$ represents a relevant point in time.

We chose a sufficiently small tie-breaker,

for nonempty instances, so maximizing evidence coverage remained our primary objective.

Mathematical lingo aside, here’s all we’re really doing:

We’re trying to preserve enough information to answer as many future questions as possible while staying within a limited memory budget.

The optimizer provides a reference point for how well a practical memory controller could theoretically perform under these assumptions.

But there’s something important I haven’t mentioned yet.

We assumed that whenever the model had all the required evidence, it would answer correctly.

At the time, that seemed like a reasonable way to isolate the memory-management problem.

As you can probably guess from the Snaggletooth example, that assumption would eventually come back to haunt me.

Separating memory failures from reasoning failures

Before going further, let me explain how I tested these memory controllers.

Imagine the following conversation:

A few days later, you ask:

Obviously, the answer is $65.

But there are actually two separate things our agent needs to do.

First, it needs to remember that the car wash cost $15 and the parking ticket cost $50.

Second, it needs to combine those two facts and calculate the answer.

The first is the controller’s responsibility. The second is the reader’s responsibility.

Now imagine the model answered $55 instead.

Was that because the controller forgot one of the numbers? Or because the language model made a mistake while adding them?

If we want to evaluate memory management, we shouldn’t confuse those two failures.

So, for the early synthetic experiments, I used a scripted reader.

Instead of asking a real language model to solve the question, this program simply checked whether all the required information was available.

If it was, we counted the question as successfully answered.

This let us study memory allocation without worrying about the language model making mistakes.

And with that, we had a fairly clean experimental setup for testing different memory controllers.

2. Teaching a memory controller what to regret

Now, our hindsight optimizer had a pretty unfair advantage.

It knew the future.

A real agent obviously doesn’t have that luxury. It has to make decisions about what to remember without knowing what someone might ask tomorrow.

So the next question was whether we could train a controller to make those decisions intelligently.

One approach I explored involved learning from hindsight regret.

The intuition is simple.

Imagine you deleted a message today, only to realize next week that you desperately needed it.

You’d probably regret deleting it.

But if that message never became useful again, deleting it wouldn’t have mattered.

We tried to formalize that idea by assigning costs to memory-management actions:

Don’t worry too much about the notation.

The indicator function is simply one if the memory item will be needed again, and zero otherwise.

The parameter $\rho$ represents the risk of failing to retrieve an archived item.

For the initial experiment, we assumed archiving was free and retrieval was perfect. That gave us:

Essentially, deleting something important is bad. Archiving it is safe.

But what about something that will never be useful again?

Well, both actions are equally acceptable!

And believe it or not, that tiny distinction turned out to matter quite a bit.

A small change that made a big difference

Our earlier hindsight teacher preferred deleting information that would never be needed again.

That makes perfect sense if you know the future.

Unfortunately, the learner didn’t know the future.

It had to guess which pieces of information would become useful, and it couldn’t always make that distinction correctly.

So while it was trying to imitate the teacher, it could learn to delete information without actually learning when deletion was safe.

Not exactly what we wanted.

We changed the teacher to accept multiple equally good actions rather than forcing it to prefer one.

Our imitation-learning loss became:

Here, $A^{\ast }(s)$ represents the set of actions the teacher considers acceptable in state $s$, while ${\pi }_{\theta }(a\mid s)$ is the learner’s probability of choosing action $a$.

Instead of insisting that the controller copy one particular action, we rewarded it for assigning probability to any acceptable action. We accumulated this loss across removal decisions and also trained retrieval predictions.

And the result?

At a 2% active-memory budget, average success on our synthetic recall task increased from 0.335 to 0.851, across three training seeds.

That’s a pretty substantial improvement from changing how we defined the teacher’s preferences!

But I want to emphasize something.

These were synthetic experiments with a free archive, perfect retrieval, and a scripted reader. The result showed that the revised teacher worked better under those particular conditions.

It did not mean we’d suddenly solved memory management for real language-model agents.

In fact, once I started testing with actual language models, things became a lot more complicated.

3. What happens when we introduce reinforcement learning?

Naturally, I also wanted to see whether reinforcement learning could improve the controller.

The idea was to reward the controller for completing tasks successfully while penalizing undesirable behavior.

In the early experiments, we used:

The controller received a task reward, but lost points whenever it took invalid actions or forced the system to evict information to stay within its memory budget.

We then optimized the discounted return:

In simpler terms, we wanted the controller to make decisions that would lead to better outcomes over time.

These rewards applied to the early controller experiments; later experiments involving evidence selection and memory writing used different objectives.

So, did reinforcement learning beat imitation learning?

Well... no.

At least not in the settings we tested.

Across those experiments, reinforcement learning didn’t outperform imitation.

That was disappointing, but it was also useful information. It suggested that simply introducing a more sophisticated training algorithm wasn’t necessarily going to solve the problem.

Moving from synthetic tasks to real conversations

Up until this point, we’d largely been testing memory management in a controlled environment.

But real conversations are messy.

People refer to things indirectly, mention important facts across different messages, change their minds, and ask questions that require combining several pieces of information.

So I moved toward benchmarks such as LongMemEval^[2], where the model has to answer questions using information from longer conversation histories.

And instead of using a scripted reader, we now had an actual language model generating answers.

This was important because it brought us closer to how real AI agents operate.

One of the clearest positive results came from query-aware memory selection.

Instead of deciding what to keep without knowing the question, we selected relevant evidence once the question was available.

In the LongMemEval configurations we tested, this improved answer accuracy by roughly ten percentage points compared with a recency-based baseline, while using about one-third as many prompt tokens.

That’s pretty exciting!

It suggested that intelligently selecting the information supplied to a model could improve performance while reducing context usage.

However, there’s another result worth mentioning.

Our neural selector didn’t demonstrate an advantage over logistic regression when both used the same features.

Yes, logistic regression. The same relatively simple model you probably encountered in an introductory machine-learning class.

I found that rather amusing.

We could build a complicated neural controller, but a much simpler model could perform comparably on the tested selection task.

So the evidence supported the value of selecting relevant information. It didn’t provide particularly strong evidence that our more elaborate controller architecture was responsible for the improvement.

But something else was becoming apparent.

Even when we supplied the model with the correct information, it didn’t always answer correctly.

The memory controller could do its job, and the language model could still fail.

And that brought me back to the strange example I showed you at the beginning.

4. The model had the answer. So why did it say unknown?

Remember our little Snaggletooth example?

Let’s look at it again.

This example fascinated me because the model wasn’t merely overlooking a relevant memory entry.

It had actually reproduced the information we needed in its own note!

So I tried a small experiment.

What if we preserved the model’s note exactly as it was, but changed the instruction for its final answer?

We added:

And suddenly, the model responded:

There it is!

The model could produce the correct answer when we changed the instruction.

Now, it’s tempting to conclude that the model knew the answer all along and simply refused to say it.

But I’d be careful with that interpretation.

What we’ve actually shown is that the model’s response was sensitive to its instructions. We haven’t established exactly what happened internally.

Perhaps the model had enough information but was overly inclined to abstain. Perhaps the prompt influenced how it evaluated the evidence.

We’d need more experiments to distinguish those possibilities.

And this becomes even more interesting when we look at another example.

What about questions that require combining information?

This time, the question was:

The model received two relevant memories:

Obviously, we’re expecting $65.

But here’s what the model produced:

Again!?

This time, though, there’s an important difference.

In the Snaggletooth example, the answer was already explicitly present in the model’s note.

Here, the note contains the two numbers, but the model still needs to perform an addition.

We then tried:

The second condition did not preserve the successful extraction of $15 and $50 in the original note. It let the model produce a new response from the beginning.

The important takeaway here is that different interventions led to different outcomes.

Simply telling the model not to answer “unknown” wasn’t a reliable way of recovering correct answers.

And we shouldn’t treat these two examples as evidence of exactly the same underlying problem.

In the first example, the model needed to extract information that was already explicitly written in its note.

In the second, it needed to combine information and perform a calculation.

Those are different computational requirements.

The examples also came from one model and prompt family, so they don’t tell us how frequently either failure occurs across language models.

But they raised a question I found much more interesting than the one I’d originally started with:

If the information is available, what actually happens between reading that information and producing an answer?

And more importantly, can we observe that process?

5. Looking inside a language model

This is where my project started drifting toward mechanistic interpretability.

Up until now, we’d mostly been evaluating models by looking at their inputs and outputs.

We supplied some information, asked a question, and checked whether the answer was correct.

But a language model is doing a whole bunch of computation between those two points.

The question is whether we can inspect that computation and learn something useful about how the model arrives at its answer.

One technique I started exploring was called the logit lens^[4].

The logit lens: Peeking inside the network

A transformer processes information through multiple layers.

At each layer, the model maintains internal representations of the tokens it’s processing. These representations evolve as the information moves through the network.

Normally, we only decode the final representation into a prediction over the model’s vocabulary.

But what if we took the representations from an earlier layer and tried decoding those instead?

That’s essentially the intuition behind the logit lens.

It lets us inspect intermediate activations by mapping them into the model’s output vocabulary.

For example, instead of only seeing the final answer, we might observe which words are associated with the model’s representations at different points in the network.

Pretty cool, right?

Well, there is a problem.

Representations can change across layers, and directly decoding an intermediate activation assumes that its coordinates are already compatible with the final output space.

That isn’t necessarily true.

Which brings us to another technique: the Jacobian lens.

The Jacobian lens: A different way of reading intermediate representations

The Jacobian lens, or J-lens, tries to account for how representations change as they propagate through the model.

Rather than directly decoding an intermediate activation, it applies a layer-specific transformation before mapping that activation into vocabulary scores.

Schematically, we can write:

Here, $h$ is an intermediate activation, and $J_{\ell }$ is an averaged Jacobian that maps the effect of that activation toward the final-layer representation.

The full readout also involves the model’s normal output processing, including normalization.

The basic idea is that we can get a more meaningful view of what concepts an intermediate representation might make the model verbalize.

Now, the reason I became interested in this technique was a paper by Gurnee and colleagues (2026), titled Verbalizable Representations Form a Global Workspace in Language Models^[1].

Big shoutout to my friend Aadil Jamari, who first told me about the J-lens and the J-space in passing over dinner.

They demonstrated something fascinating.

Consider this question:

To answer correctly, the model needs to identify the animal as a spider before determining that it has eight legs.

But the word spider isn’t explicitly mentioned in the prompt.

Using the Jacobian lens, the researchers observed representations associated with spider appearing in intermediate layers.

That alone is interesting, but it doesn’t necessarily prove the model was using that representation to produce its answer.

So they went a step further.

They intervened on the model’s internal representation, swapping the relevant lens coordinates for spider with those for ant.

And the model’s answer changed from eight to six.

WOW.

The intervention provides evidence that the intermediate representation played a causal role in the model’s computation.

Rather than merely observing a token that happened to appear in an intermediate readout, the researchers changed that representation and observed a corresponding change in the output.

That distinction is incredibly important.

And naturally, it made me wonder whether we could use similar techniques to investigate our memory failures.

6. Can we actually see where the model goes wrong?

My initial hope was that we could trace how evidence turns into an answer.

Imagine giving the model a memory that contains a particular fact.

As that information passes through the network, perhaps we could observe the relevant concepts appearing in intermediate representations.

Then, by comparing successful and unsuccessful answers, we might identify where the computation differed.

At least, that was the idea.

And some of our early observations seemed promising.

We occasionally found answer-related tokens appearing in lens readouts even when the model ultimately produced an incorrect answer.

At first glance, you might think that’s evidence that the model had internally recovered the answer but failed to express it.

However, there’s a problem with that interpretation.

Just because a token appears doesn’t mean the model derived it

We ran a control experiment where we removed the supporting evidence.

And interestingly, a substantial amount of that apparent answer presence survived.

In other words, some of the words we thought might indicate successful reasoning were still appearing even when the model no longer had the necessary information.

That’s a pretty important warning sign.

A model might associate certain concepts with a question because they’re statistically plausible, not because it successfully derived them from the evidence.

For example, if I ask you what animal spins webs, the word spider is an obvious association.

But if a particular answer token appears in an activation readout, that doesn’t automatically tell us how the model arrived there or whether it used the evidence we provided.

Recovering a word from an internal representation is not the same as proving that the model computed the answer.

And that made our interpretability experiments a little more complicated than I’d hoped.

What did the smaller models show?

We also evaluated the lens methods on smaller language models.

The results were mixed.

For Qwen 3B, the Jacobian lens passed our answer-rank improvement test but failed an aggregate intermediate-recovery comparison.

Other small models passed their respective evaluation gates.

So while some readouts were useful, we couldn’t simply conclude that smaller models were incapable of representing abstract intermediate concepts.

We also experimented with changing instructions and steering model activations.

Some interventions recovered correct answers, which was encouraging.

Unfortunately, they also introduced new errors.

We went on to explore activation patching^[5] and evidence substitution, trying to determine whether particular representations had evidence-specific causal effects on the final answer.

But the experiments had limitations that prevented us from cleanly distinguishing all the failure mechanisms we were considering.

Admittedly, this was frustrating.

You’d hope that looking inside the model would give you a straightforward explanation of what went wrong.

Instead, we found that the measurements themselves required careful interpretation.

And I think that’s one of the most valuable things I’ve learned from this project so far.

There are three very different questions we need to distinguish:

1. Did the model answer correctly? This measures its observable performance.
2. Can we recover an expected intermediate concept? This tells us what a particular interpretability method can reveal.
3. Does changing that representation alter the answer in the predicted way? This gives us evidence about its causal role in the computation.

These might sound similar, but they’re not interchangeable.

A model can answer correctly without exposing an expected intermediate through our lens.

A lens can expose an answer-related token without showing that the model derived it from the evidence.

And an intervention can change the output without necessarily proving our entire explanation of the original failure.

So we have to be careful about what each experiment actually tells us.

7. Where do we go from here?

When I first started this project, I was trying to answer a relatively practical question:

How can we manage an AI agent’s limited memory more effectively?

That led me to optimization, hindsight oracles, imitation learning, reinforcement learning, and eventually real conversation benchmarks.

And we did find something useful.

Selecting relevant evidence can improve answer accuracy while reducing the number of tokens supplied to the model.

But in the process, we also discovered that memory selection was only part of the problem.

Even when the evidence was available, the language model could still fail to answer correctly.

That observation led me toward a different research question:

Can we identify the intermediate computation that connects available evidence to an answer, and understand how that computation differs when the model fails?

This is the direction I’m now exploring for my undergraduate honors thesis.

Starting with simpler, controlled questions

One of the challenges with real conversation benchmarks is that they’re messy.

Questions vary in difficulty, evidence can be scattered across multiple entries, and there are often several different reasons a model might fail.

If our goal is to understand the computation itself, it might help to start with a more controlled setting.

For example, imagine we give the model:

And then ask:

To answer this, the model needs to connect Maya to Alex, then Alex to Leeds.

The answer is Leeds.

Now imagine we systematically change the names, locations, relationships, or the question itself.

We could study whether the model’s internal representations change in ways consistent with those modifications.

For example, if we replace Leeds with Bristol, does the representation associated with the intermediate result change accordingly?

And if it does, can we intervene on that representation and predictably change the model’s answer?

That’s the kind of experiment I’m interested in.

More specifically, I want to study three kinds of tasks.

Direct extraction: The answer is explicitly present in the evidence. For instance, identifying the Snaggletooth action figure.

Relational composition: The model has to connect multiple facts, like determining where Maya’s supervisor works.

Temporal updating: The model needs to determine which information remains valid after something changes. For example, someone moves from Leeds to Bristol, and the model needs to use the updated location.

By constructing matched successful and unsuccessful examples, we can start investigating where these computations differ.

We can inspect activations across layers and token positions, identify candidate intermediate representations, and then test whether intervening on them changes the final answer in predictable ways.

Of course, there’s no guarantee that the relevant computation will be easily readable through a lens.

And even if we find a promising representation, we’ll still have to establish whether it actually plays a causal role.

But that’s precisely what makes this direction interesting.

Instead of assuming that every incorrect answer is a failure to retrieve the right information, we can start investigating whether the model failed to derive the answer, represent something correctly, or use a relevant computation when generating its response.

8. The bigger picture

Looking back, I find it rather amusing how much this project has changed.

I started with an optimization problem inspired by an algorithms class.

The original goal was to decide which memories to keep, archive, or delete so that an agent could answer as many future questions as possible.

We built a hindsight optimizer, trained controllers, experimented with reinforcement learning, and compared different selection methods.

But underneath all of those experiments was one assumption:

If the model has the information it needs, it should be able to answer the question.

And then a three-billion-parameter language model looked at a note containing a Snaggletooth action figure and told me the answer was unknown.

That was the moment the original assumption became something worth investigating.

I don’t think the experiments we’ve conducted so far give us a complete explanation of why this happens.

The interpretability results are mixed, some interventions help while others introduce errors, and the selected examples don’t establish a general mechanism.

But they do give us a much clearer idea of what to investigate next.

Perhaps improving AI agents isn’t just about building better memory systems, adding more context, or developing more sophisticated controllers.

Perhaps we also need to understand what happens after the model receives the information it needs.

After all, having access to the right information and successfully using that information are two different things.

And understanding the gap between the two is what I hope to investigate next.