Titans: Neural Long-Term Memory for LLMs, with author Ali Behrouz

Today, I am thrilled to share my conversation with Ali Behrouz, a PhD student at Cornell and lead author of the fascinating paper on integrated large language model memory, “Titans: Learning to Memorize at Test Time.” This paper represents another significant step forward in addressing what I’ve often called the missing middle in memory for large language models. We’ve got world knowledge baked deeply into the model weights, and we’ve got attention-based working memory in the context window. But that missing middle layer—the integrated, persistent, and ever-evolving long-term memory that humans use to maintain coherence and identity over time—still seems like a necessary piece for success in high-context environments, and has been a frontier to watch in AI for some time now.

It was once famously said of computers that you could see them everywhere but in the productivity statistics, and to a lesser degree, that’s been true of AI as well, at least so far. In part, this is because the technology itself is really only recently useful, and people take time to adjust. Tyler Cowen recently went reverse Oprah, pointing out to an influential audience, “You are a bottleneck. You are a bottleneck.” But I think often the bigger barrier is best summarized by another famous Tyler quote: “Context is that which is scarce.”

To be honest, I’ve never been entirely sure what Tyler meant by that in the context of humans, but in the context of AIs, it makes perfect sense. Today’s AIs know a literally superhuman amount about the world at large. But out of the box, they know little to nothing about the individuals and businesses that they’re meant to serve. Meanwhile, assembling and maintaining relevant context for them—especially because it’s often spread out across Slack, email, documents, GitHub, meeting transcripts, task management systems, you name it—is tedious work. And the AIs have really only recently, I would say truly with Gemini 2.5 Pro, started to properly reward it.

But for the sake of argument, imagine a world in which context is not scarce for AIs—a world in which an LLM trained specifically for a major company, say GE or 3M, knows as much about that company, including its products, its history, its team, its internal processes and debates, and its finances, as it does about the world at large. Obviously, such a model could be created. No company has more than 1% of the data on which the AIs are already trained, and if it were created, it would immediately know more about the company than any single person at the company.

It might still be trained to search official records to ground its analysis and work. But unlike today’s models, it would know, in a way very similar to how humans do, when it’s actually found what it’s looking for. And it would almost certainly pick up many of the subtle patterns that constitute what I call “how we do things around here,” which could make it relatively easy to manage, especially compared to the overall process of hiring, onboarding, and retaining human knowledge workers.

Startup costs for a model like this could reach into the millions or perhaps even tens of millions of dollars, which is notably roughly where OpenAI has publicly priced its custom models offering. But for a 100-year-old enterprise, that would still be a bargain for an AI that you can drop in and have do a significant portion of the work at the company, particularly since you’ll be able to amortize those costs across as many copies as you need. For smaller businesses, meanwhile, which have much less data, using today’s fine-tuning prices as an anchor, I would expect costs to be more like tens of thousands, maybe into the low hundreds of thousands of dollars. Obviously, still affordable.

We’re starting to get a glimpse of this future as individual ChatGPT users get new memory features these days, but I’m not aware of a productized version of this that works well at scale. And I think that’s ultimately because there’s still a gap in the foundation models themselves, which scaffolding isn’t quite enough to fix. The bottom line, then, is that it seems plausible to me that the main thing between where we are today and a future full of drop-in knowledge workers that begin to very quickly and dramatically disrupt the labor market is simply a breakthrough in long-term memory. And that’s why I think today’s topic, Titans, is such a big deal.

Unlike RAG-type systems, which store data explicitly and then make it searchable via a mix of traditional, embedding-similarity, or graph searches, or even more integrated strategies like Mamba and other state-space models, which encode memory as a matrix of numbers that gets updated as part of each forward pass, the Titans architecture that Ali and his co-authors propose uses a neural network, which itself is updated via gradient descent at runtime, as the LLM’s memory module. This is a qualitatively different approach and, in my humble opinion, represents a significant conceptual advance.

So with all that in mind as motivation, in this conversation Ali and I explore the technical details of Titans, including how he conceptualizes and takes inspiration from human memory systems, how the associative memory loss function works, the role of surprise and momentum in updating the neural memory module, and the various approaches they experimented with for integrating this long-term memory module with the standard attention mechanism.

Beyond the technical details, this conversation also offers a fascinating glimpse into how a highly original and obviously quite brilliant researcher thinks about pushing forward in a relatively new architectural direction. I was really struck by how many times in this conversation I wanted to dig in and understand the reasons behind the decisions that Ali and his co-authors had made. But his response was basically that they see all of this as very early foundational work, and so they just did the simplest thing possible for now, trusting and expecting that others will come along to improve and refine it later.

Overall, I really enjoyed this conversation. While it is pretty technical at times, I think Titans is one of those relatively few papers that is worth taking some time to grow into. The intuitions behind it are elegant, and this line of work might just produce the last major technical unlock needed for AI to hit an inflection point in economic value and impact.

Ali Behrouz, a PhD student at Cornell and lead author of the fascinating new paper on large language model memory, Titans. Welcome.

Thank you very much. Thanks for having me.

I am really excited about this. I think every once in a while, one of these papers comes along that just has me thinking and thinking and thinking, and this is one of them. So I am really excited to get into it.

I guess, just for a little background, I’ve kind of had a little obsession with the memory aspect of large language models for a while now. So when I saw the Mamba paper drop roughly a year before the Titans paper, I thought, “Boy, this is really a big deal,” because I think everybody kind of knows that large language models have this incredible world knowledge that’s super vast, but it’s fixed, right? The knowledge cutoff date can be somewhat recent or it can be more in the past, and whatever it is, that’s sort of a crystallized thing that doesn’t really change too much between versions of a model.

Then, of course, we’ve got the runtime memory that is the context window, and the attention mechanism is great for connecting all the different tokens to each other token and making sense of what’s going on there. But I’ve often called the gap between those the missing middle in memory. You can fill the context window pretty quickly, but it’s proven really hard for the field to create a sort of sustained coherence where a model can know who it is, know what it’s done, and know where it’s trying to go.

I think that’s going to be a really important piece of the overall puzzle as we think about trying to create effective agents that might be long-running in the real world. So Mamba was a key moment where I thought, “Oh, that seems like a major step forward.” And this is, with a different approach, another one of these things that I think really shines a lot of light on the path.

With that preface, tell us just a little bit about your background and how you got interested in this work. What angle are you coming to it from? I think the paper is really interesting in that it’s got some elegant math, clear analysis of the recent history of people’s attempts to bridge this gap, and even a little bit touching on human memory that takes inspiration from it in some ways. I’d love to hear your inspiration for this work, maybe for starters.

Sure. Honestly, one thing that I need to say is that my path in this direction is very unusual. I started doing research during my master’s, mainly, and I started working on some graph algorithms and these kinds of things—completely, very far from this area.

When I did some research, I realized that some of the tasks I was trying to do were much easier when we used graph neural networks and these kinds of things. So I started learning more about graph neural networks to do those tasks. I mainly focused on questions about anomaly detection and these kinds of tasks in graph learning. At the end of my master's, I realized that there were some interesting connections between things like anomaly detection in graph algorithms and graph datasets, and also some neuroscience tasks.

For example, let's say we want to detect a disease, or we want to understand some disorders—brain disorders, for example—or, generally, these kinds of tasks. That actually made me more interested in approaches like neuroscience-inspired methods for deep learning. My intuition is that humans are very effective—I mean, the learning process in humans is very effective and very efficient. We can learn with a small number of samples, and that's a result of millions of years of evolution.

So it's really hard to believe that we are smart enough to beat millions of years of evolution and come up with an architecture that is even more efficient and effective than humans without passing through those steps. We can come up with an architecture that is very similar to humans and then improve it to, for example, build that superintelligence—whatever we want to call it—but I really believe that before that, at least at this time, we need to get inspired by how our brain works.

Actually, the start of my PhD somehow coincided with so many great papers about sequence modeling and alternative architectures. I think at that time there was the S5 paper, which was one of the first papers introducing state-space models, and there were also RetNets and all these other great models. All of them motivated me to try to understand what we can do with alternative architectures rather than transformers.

Personally, I think one thing that's different is that there are several people working on alternative architectures, and each of them has their own perspective. I believe all of them are really good and great, but there are some people who don't believe in attention and want to fully replace it with RNNs. There are other people who believe in hybrid models, and I'm one of the people who believe in hybrid models because I really think the attention part is necessary for accurately modeling dependencies between different tokens in the context that we have.

Yeah, I'm with you on the instinct that hybrid models are ultimately the way to go. So, a quick review of what the problem is with the attention mechanism: It works great, but the memory requirements get longer and longer as you keep extending the sequence, because you're doing this all-token-to-all-token calculation. So the size of that memory footprint grows as the square of the sequence length. I think everybody watching this probably already knows that pretty well; just in case, that's the core problem that we're trying to solve.

So that can't go on forever. Now, I guess you could have—and we've seen schemes where it's not necessarily part of the model itself—some arbitrary disk space, but arbitrary disk space is quite different from arbitrary in-memory computation. So intuitively, it feels like we need something more like the human brain, which is finite in size. Our brains aren't growing with every time step throughout our lives, and we need some elegant updating mechanism for memory that allows us to keep what's important, while knowing that it is a finite state. That means we also have to let go of some things over time.

Intuitively, we all know that our memories do this, but we haven't quite cracked the perfect way to do that in the context of an AI system. I wanted to take just a couple of steps back for a little more foundation. One would be to get your sense of how you conceptualize the human system, maybe get a few thoughts on how you understand linear approximations of attention, and then a little bit on how you understand state space—or, perhaps a little more generally, how you understand prior attempts to create these finite-size, constant-update-time architectures like Mamba and RetNet and others that have come before.

Do you want to take those one by one? It doesn't need to be a full lecture, obviously, but if there's anything that you feel is distinctive about your perspective on those 3 topics, I'd love to hear how you conceive of them and how they motivate this work.

Let me start with some explanation about traditional RNNs. Let's look at what RNNs are doing at a very high level. Basically, the data comes in, we project the input, or the data, into the hidden space, and we also project our memory, which is a vector. Usually, people in the literature call it the hidden state. We project that, add these 2 together, update the memory, and repeat this process.

Anytime we want to extract something from the memory, we just multiply the hidden state, or our memory, by a matrix—usually H × C or C × S, depending on how you define the dimensionality—and that equals the output, which is Y_t. That's how a linear RNN, or generally an RNN with nonlinearity—we can simply apply that here—works.

Now let's go to the next model. Linear attention in the causal setting can again be written as a recurrent neural network. Again, everything is similar: Data comes in, we project the data, there is a hidden state, and we update the hidden state, and so on. But now this hidden state is a matrix-valued memory.

If we want to somehow say that there are differences between a linear RNN, a traditional RNN, and, on the other hand, a linear transformer, one thing we can see is that traditional RNNs and modern linear RNNs—for example, Mamba—use vector-valued memory. On the other hand, linear attention uses matrix-valued memory.

Now let's say we know all these things. The question is: What is memory in our brain? Is it something like a set of neurons that we define, or is it something like a larger neural network of interconnected neurons? I think it's the second one. We cannot say that there are 3 different neurons encoding all the memories and all the stuff and everything.

Usually, when I talk about Titans, I start my talk with some explanation about different perspectives. I think here it might be useful to start from that again. There are different perspectives that, over time, actually help us design new architectures. When we're thinking about LSTM and that kind of model, they're inspired by the brain. LSTM stands for long short-term memory.

We wanted to design something very similar to our brain, but the tools we had at that time, our understanding of the brain, and all these things change over time.

So basically, that’s the main difference between the traditional perspective on human memory and the current perspective. Another thing, as time passes, is the transformer era, in which most of the focus is on efficient matrix multiplication. We have a lot of models designed based on, for example, efficient matrix multiplication algorithms, like MLP-Mixer and its variants, structured matrices, and all these things. So basically, that’s another perspective that we have, and each of these perspectives gives us some new intuition, some new paths that we can follow to find new architectures.

Probably starting from 2020, we have the perspective of dynamical systems—I mean, treating the sequence as a dynamical system and using state-space models for modeling sequences. So again, we have some new intuition and new insights about how we can design efficient and effective sequence models. I really think that now we need to go farther and go back to the human memory perspective, but with some novel understanding of how human memory works and how we can incorporate some novel techniques that we are learning.

In this perspective, the question is: What are we missing in previous architectures that we need to address? The first thing is that our memory is not something that we can break into pre-training and post-training; it’s constantly learning. That’s a very important part, and it’s very good, for example, for RNNs, because now, from the TTT paper, we know that most of these RNNs—or all of them—are doing test-time training. So basically, it seems that we are in good shape with modern RNNs, but probably we need to do additional things to make them more effective in test-time training.

Another thing that we might miss in this perspective is, as I mentioned, the shape of the memory. Is it a vector? Is it a matrix? Or is it a neural network that might have different architectures, different designs, and all these things? I believe in the last one. I don’t think that, for example, it’s suitable to say we could design architectures with so much data and scale them, and all of them are done with only one vector as the memory of your model. So that’s somehow an oversimplification of this design. Probably we need to take more complicated architectures as our memory.

Another thing that we need to consider in this perspective is that, let’s say, we have some recurrent neural networks. These recurrent neural networks have a memory that is fading over time. On the other hand, we have attention, which considers all the pairwise interactions between the tokens within the context window.

That’s very similar to our short-term and long-term memory because it’s really hard, or even impossible, to hallucinate about information that we just got in the past 30 or 20 seconds. We really have all the information in our memory. We don’t hallucinate about any details of a piece of information that we just got. But it’s very likely that we hallucinate about information or an event from 20 years ago. There are so many details that we don’t remember, and there are so many details that have even changed. For example, we remember something that wasn’t true at that time. We’re just hallucinating about that.

If we think about an RNN as a fading memory, we can see that this decay mechanism—this inherent decaying mechanism, not the forget-gate part—exists because, in any design, RNNs have this decaying mechanism. We’re always adding new data into the RNN, and we might ignore some of the data that we have. So this RNN is very similar to our long-term memory. It has the ability to manage information and all these things. On the other hand, we have the attention part, which is very accurate and very similar to our short-term memory.

I think we need to use the combination of these two to design more powerful architectures. But the question here is: How can we do that? Even in the neuroscience literature, there are some controversies about how we pass memories from short-term to long-term and, for example, how we do these things. So basically, that’s a part that we don’t have anything to get inspired from. That’s a challenging part and probably an important research path in the future, because we might design some models that are more effective in passing memories from short-term to long-term and all these things.

So basically, that’s the architecture that I believe in. Generally, this kind of thinking led us to the design of Titans.

Cool. So I guess one thing to really emphasize is that, in some sense, the most fundamental change here with this work, compared to everything previously that I’ve seen, is that the memory module in Titans is itself a neural network.

We’re kind of moving from—and you said this, but just to really hammer the point home—the earlier recurrent architectures had either a vector or a 2D matrix or whatever, but basically just a bunch of numbers that information would be projected into and then also sort of projected out from at each inference-time step.

But within that, you just have these numbers, and they don’t interact with each other. There’s not really much going on there, right? It really is just sort of a place to store the results of these projections in and projections out. Moving to a neural network as the memory module—and you use an MLP—and I definitely noted the part in the paper where you said this opens up a whole new research direction in terms of what the best possible architecture for that submodule is.

An MLP may be hard to beat, but it’s probably beatable. We’ve certainly seen many other architectures be better for many other purposes, but moving to this architecture creates a dynamism within the memory module and creates the potential for information to interact within that piece, which is qualitatively different and seems, as you said, much closer to the way that our own memory systems work.

One thing I don't have a great sense for, and I'm sometimes surprised by, is how much of this is all reformulatable if you're smart enough about the math, such that things can become more equivalent. This is not intuitively obvious to me, but the linear approximations of attention you noted can be reframed as a recurrent neural network. One of the big things we've seen through this whole progression, from the linear approximations and then to RetNet and Mamba-2, and many others besides along the way, is a much more granular management of that internal state, and also, with Mamba in particular, a jump to input-dependent ways of managing that internal state.

Early on, you would just have, okay, here's this data: a fixed sort of kernel that's going to do this projection, and no matter what information comes in, we always compress it in the same way. That kind of is what it is. With Mamba in particular, you had an input-dependent and highly granular way of updating. So, if I recall correctly, each number in that matrix could be updated with a different strength. You weren't just doing one kind of overarching, single decision for how much to decay and how much to emphasize the new information; you were doing it in a much more granular way. And again, that depended on the inputs. So that was quite interesting.

Another element of this that also always comes up is how much of this is dictated by first principles and how much of it is dictated by the hardware that we have available to run it. Mamba versus Mamba-2, I thought, was a really interesting illustration of that, where Mamba-2 was actually a less granular method for managing that internal state. It was a coarser update function. You'd think, well, geez, how do they get a better model out of a less granular, or more coarse, way of managing the memory?

The answer, as I understand it, is basically that by making that sacrifice of the super-granular Mamba-1 structure, they were able to make the whole thing run a lot faster. So, for a certain number of GPUs, they were able to train a lot more, and so you get a better final model, although it does have, in some sense, this coarser internal structure.

To put that in the form of a question: How much of all of this do you see as being very dictated by the hardware that's available? Is there a fundamental break here when we move from a matrix-shaped memory module to a network as the memory module? Or is there some fundamental equivalence, perhaps, as well, where this could all be, with enough insightful math, understood as being in some sense the same thing?

So let me use the last part—I mean, the last page of the Titans paper. We did some comparison between, for example, Titans and some recent modern architectures like DeltaNet or, for example, TTT, Longhorn, and all these models. What we did in that part was to say how we can connect, for example, DeltaNet to Titans, what is changing in that sense, and, now let's say we have the perspective of TTT, what is changing in that sense, and so on and so forth.

One thing that we need to consider is that all these models are connected. There are some mathematical formulations that we can use, I mean, at least in the sense of this test-time training framework, that unify all these things. Basically, all of them are connected, and there are some small changes in each of them, but I don't see the contribution from that side. I really believe that the value of the contribution is to show how you want to do some future work, because probably no architecture is endgame, or nothing is endgame. We are progressing over time.

It's not good, even for science, to say that this model is coming and it's the end of the world. We are not making any progress over time. Even transformers are really great. Even when we have transformers, we have better implementations of transformers, like FlashAttention, and all these things. We have better additional components to transformers that make them Transformer++ and all these things.

The important part, I really believe, is not just to deliver a model, because probably there are a lot of things that we can do to make a model better, but to deliver a new perspective, to deliver something that leads to future studies. Personally, if I want to say what the impact of one work is, for example, from a very old time, I would say that the impact is something that shows the future work, and those future studies led us to the state that we have right now.

So I think that's a really important part. Yes, we can look at all these models that are very connected, but most of them are from different perspectives, and each perspective can lead us to different types of thinking about how we want to improve current architectures and all these things. The question is, we don't actually know the answer to this question, but which of these directions will, for example, lead to a good architecture that is better than other methods that we know at that time?

For example, let's say that we use a matrix-valued memory in Titans and, for example, let's say we remove some layer normalization that we're doing inside the memory, which makes the memory linear. But let's say that we don't do all those things. Basically, in that case, our approach can be very similar to, for example, Gated DeltaNet or DeltaNet. If you look at these things, you can say that there are all small differences between these models.

For example, as I mentioned, you can remove some of the components that we have and also consider the special case of our approach, which uses a simple matrix format for the memory. Then there are some connections to these models, similarly for TTT, similarly for Longhorn, and all these models. But even a small change in these kinds of approaches can lead to completely different architectures, in terms of the perspective that they can bring and all these things.

It's really important to see what the connection is between these models and also which types of approaches can help us to design future architectures. In this direction, I can say that the focus is that we have a fixed-size memory, and this fixed size is something like lifelong memory for the model. We are adding data into the memory, and so the question here is which types of approaches we want to take.

One approach is to make memory management better over time. What does memory management here mean? It's just the recurrent formulation that we have. Different recurrent formulations help us to better manage the memory and better understand which information is worth memorizing or which information we just need to ignore. That's the important part, in my opinion.

But there's another perspective here that says we need to do some things to have better memory management. On the other hand, we need to see how we can use different architectures for our memory to make it more powerful. So we started from a vector, then went to a linear layer or a matrix, and now we can have an MLP with 2, 3, or 4 layers. So what's next? Do we want to use convolutional networks inside the memory? Do we want to make deeper architectures in the memory and all these things?

I really think that these 2 different approaches are very separate. It's very hard to connect these 2 because they are answering different questions, and we need to see which one is more promising. Definitely, we can use both in future architectures. We can have better memory management with better memory architectures.

But one thing that I want to emphasize is that these 2 paths are separate, and they're trying to answer different questions in this domain. That's, I think, very important.

Cool. Well, let's get a little bit deeper, then, into just the nature of the architecture itself. There are a lot of little, detailed choices that you've made that I understand are not the end of history. I always say transformers are not the end of history, and I totally understand that this paper represents more the beginning, probably, of a new direction, certainly not the end.

If I could try to summarize the overall architecture simply, it's akin to a transformer in the sense that you still have attention as a pretty core mechanism. But now, instead of being truly all the tokens that the model has ever seen in the sequence, we now have essentially a sliding attention window, which is something that we've seen in the past with various schemes.

Then there's also this bit that I didn't have a great sense for, which is the persistent memory—I think that's how it's called in the paper—and it's also called learnable data-independent weights. What I see in the diagrams is that there's this sort of persistent memory layer that gets seemingly always put at the beginning of the sequence at every time step. So you sort of have the persistent memory, followed by the output from the long-term memory module, which we'll describe in more detail in a minute, and then followed by the current sequence, which is the sliding attention window mask.

Then all of that kind of goes into the attention mechanism and gets processed through the larger supernetwork, in more or less a transformer-like way. Tell me if I'm missing anything important there. I didn't quite have any great intuition for why you thought that the persistent memory—where did that idea come from? I have seen some things in the past.

I remember a paper called Attention Sinks, or something along those lines, where it seemed like they had found that there was basically an overdependence on the initial tokens. The initial tokens were too important in some ways. Having these extra early tokens that might not be very important, but also not allowing the model to overweight the actual runtime's first tokens, seems to be a performance benefit.

I don't know—deconfuse me on that. Where does that persistent memory bit come from, and what role should we understand it as playing?

The persistent memory part might not necessarily be needed, but it completes our architecture. If I want to explain why we are using that persistent memory, there are different reasons for it. The first reason is that we motivate the paper by asking how we can design something similar to human memory. Basically, there are some long-term memories, some short-term memories, and also something like persistent memory that encodes knowledge about the task we are doing.

It's not related to the data; it's about general knowledge of the task. That makes the design complete, so we put it there. In practice, there are some improvements: when we use that persistent memory, we get a very, very slight improvement in the architecture. There is no harm in that, so that's the main reason from a technical perspective that we are doing these things.

On the other hand, another reason we use persistent memory is that we concatenate the long-term memory with the current context. As you mentioned, attention can focus more on the initial tokens. That's challenging and might result in a drop in performance. Having some data-independent, learnable parameters at the beginning of the sequence can help us improve performance.

One thing I need to say is that the idea of adding additional persistent memory to the sequence comes from an earlier paper than Attention Sinks, which we already discussed in the paper. Basically, what that paper says is that we have the attention mechanism in transformers, and after that we have the MLP. One way to view the MLP is as an attention mechanism.

Let's say that we have a 2-layer MLP. The formulation would be something like W1 × W2 × X, with a nonlinearity between W2 and W1. If you look at that, you can see that W × X is very similar to, for example, Q × K in transformers or in attention. Then another W is basically our value matrix.

That's very similar to the attention mechanism. But what is the main difference? The main difference comes from the nonlinearity. We usually use, for example, GELU in the MLP, or other nonlinearities, whereas in attention it's softmax. What they say is that we can use softmax here. In that case, what we are doing is concatenating some additional learnable parameters to the beginning of the sequence. When we apply attention on top of that, it seems that we are already applying the MLP part, so we don't need that part anymore. That's the main intuition from that paper.

Again, as I mentioned, there are different perspectives. There is one mathematical formulation, but there are different perspectives that all lead us to this mathematical formulation. That's generally the main motivation for the persistent memory.

I understand that some people might not use that persistent memory when they want to use Titans, and that's understandable because it can save us some parameters and is simpler in design. On the other hand, we can gain a slight improvement in performance.

Gotcha. Yeah, that's really interesting. In a way, that is also why I have been so fascinated with studying AI broadly: there are just so many different angles on it, so many different perspectives. It feels like, in some ways, we're all trying to get at the same ground truth, but the different perspectives do have quite different value to bring, even if there is some underlying mathematical equivalence.

If I take a big part of what you're saying, it's that the different perspectives get people thinking about different future directions in ways that underlying mathematical equivalences may not actually be so useful for. That's definitely true for me.

Okay, that's interesting. Let's talk about the long-term memory module. Obviously, this is again the core thing, right? We are now moving into the world of having a network that is the memory. I don't know if you want to take a minute and talk about the read or, sort of, write-and-read paradigm. We kind of already covered that, but is there more to say about the modern paradigm of writing to and reading from memory before we get into the specific details of what you've implemented here?

Basically, the intuition here is very similar to what I mentioned about the simple paradigm of RNNs and linear attention, and these kinds of things. As I mentioned, data comes in, we project the data into a hidden space, project the memory, consider their summation, and update the memory. We repeat this process, and so on and so forth.

Generally, what we are doing here is writing to our memory. The data comes in, and we ask how we want to write this data into our memory. The way we do it in a linear RNN is to project it and add it to the memory.

There is another part. Let's say that we have this memory: how do we want to extract information from it? We can think about sending a query to the memory and asking it to give us the corresponding information. We can model that by multiplying our query by the memory, or passing it to the memory.

If your memory is a simple vector, you can simply multiply by it. If it is a simple linear layer, you can just multiply your query by the memory. On the other hand, if your memory is a neural architecture, reading from memory becomes something like a forward pass. You need to pass your data into your memory, and the output of your memory is the corresponding information about your input. That's how this write-and-read intuition works.

Cool. Let's get into the details of this, because I found it really interesting. I have to give a shout-out to Grok 3 for helping me work through some of the math and develop my intuition for it. Grok 3 has definitely been shown to have a number of interesting properties and some interesting issues, but it does perform quite well when given a paper like this and asked conceptual questions about it.

One fundamental decision that you made here is to create memory that you describe as associative memory. Basically, the memory module is an MLP, right? It can take some input and give you some output. The loss that you're minimizing there is the difference between what the memory outputs for a given input's key vector and the value that same input actually ends up with on the other side of the attention mechanism.

I interpreted this as an approximation of the attention mechanism. What we're doing is saying, "We want this finite MLP, given the key value for some input, to be able to predict what the value output would be for that same input." Then we'll do our runtime updates of the MLP.

It's probably hopefully obvious to people at this point, but one major change here from previous architectures is that this MLP is actually undergoing a gradient-descent process at runtime. The weights of the MLP are changing. In previous architectures, you had, say, in Mamba, this matrix of numbers that changes with every time step. Here, we have an MLP that changes with every time step, and it changes through gradient descent, as it usually does.

The loss function that defines those gradients is this prediction: we want the memory, given a new input and the keys for that input, to be able to predict the values for that input. Maybe you could take a second and talk about how you think about the queries, keys, and values structure of attention.

People have encountered that in the past. The simple shorthand I can recite procedurally is that the query portion is what a given token is looking for, the keys are what information it has, and the values are the payload of what information then gets fed into the rest of the network.

I've always held that understanding relatively loosely because these architectures are weird. Interpretability is an immature science, and it shows us a lot of weird things. Even if that was the idea people had in their heads when they set up this architecture, does that mean it's really working that way?

But you seem to really embrace that intuition, or that understanding of what's going on, and reuse it in the design of this memory module. Tell me how you understand queries, keys, and values, and tell me if I'm right about how you applied that to the design of this.

I think one way we can describe the attention mechanism, as you mentioned, is that we can view attention as an associative memory. Basically, we have keys and values.

They're connected. We want to pass keys into the memory, and the memory is responsible for finding the value corresponding to that key and passing it as the output. That's how the memory should work. The main difference here is that we might not have the exact value of the key when we are doing inference and these kinds of things. That's where we use a query.

So we have some stored pairs of keys and values in our memory. We want to send a query to the memory, and the memory, as I mentioned, is responsible for finding relevant information for the query and passing it as the output. How should we do that? Let's say our memory has K₁ to K_L as our keys, and we have another query that is generally different from K₁ to K_L. How can we find the relevant information from this memory?

One way to do that is to see how similar this query is to the keys that we already have in the memory, and how we can describe this query as a combination of the keys that we have in the memory. The way we can do that is by using matrix multiplication, because the dot product can help us understand the similarities between each pair of vectors. We have a query, and we consider the dot product of this query with the keys to understand how we can describe this query using the keys that we already have in the memory.

Based on that similarity, we can extract information from the memory, because we already know the connection between keys and values. When we can describe this query as a combination of some keys, we can find the output of this query—the corresponding output—as a combination of the values that we already have in the memory.

That's exactly what we're doing in attention. We consider the multiplication of Q and K to find the similarities, and then multiply that with the value to extract the information from the memory. That's one way to think about attention.

But if we write attention as it's done in the Test-Time Training (TTT) paper, one thing that we can see is that attention is the nonparametric solution to the loss function that we have. Basically, these keys and queries are interconnected and assigned to each other; attention is the nonparametric solution to that assignment. When we're talking about an RNN, we're talking about applying an optimization algorithm—gradient descent, or gradient descent with momentum—on top of this. Basically, an RNN might result in weaker performance because the attention part is the nonparametric solution to that problem. On the other hand, we have some efficiency gains that we can use.

Let's just dwell on this for a little bit longer, because I feel like this is maybe the most important thing for people to develop an intuition for. If they can grasp this, a lot of other things will make intuitive sense relatively naturally.

We want to have a memory module that says, "For this given input, feed it into the memory module and get the most relevant historical information through all of time." To do that, we have the query aspect of a given input, which is understood to be—or was maybe originally conceptualized as—what that token is looking for. We can feed that in, and we want it to come out with all the right information.

How do we update that memory module over time so that it contains that information, so that the query of the latest input at runtime actually gets the stuff that we need? At every previous step, we need to have encoded it. How do we encode it? We say, "Well, let's take the keys of all those previous steps and train this memory module to output the values associated with those same steps."

This builds on the fact that, in the attention mechanism, it's the similarity—the dot product of the new token's query versus all the previous tokens' keys—that determines what portions of the values will actually be used in the downstream calculations. At every runtime step, we want to make sure that, given something similar to the keys of this input, we can output the value so that we have that payload information to pass into the rest of the network.

We update it so as to be able to do that in anticipation of actually getting the query that says, "This is what I'm looking for." We want to store, essentially—we want to get the network to be able to predict what each token has. That's how we're updating it as we go, so that we can later say, when a new token is looking for what those previous tokens had, that we can return something similar, at least, to the values that would have been produced had we done the full explicit attention mechanism. Anything wrong about that?

No, I think that's the correct way of thinking about this mechanism.

Okay, cool. I think that's really super interesting. It took me a little while to get there, but it was worth the work. Again, I appreciate Grok for helping me work through some of the notation to make sure I was understanding it correctly.

I think this will probably also help a lot of people understand the attention mechanism itself a little bit better. It gives me a little bit more confidence, too. The fact that this all works gives me a little bit more confidence in the idea that the attention mechanism is maybe actually working the way people have described it as working, because it's really easy to fool yourself into thinking that you know what's going on in these massive computations and number-crunching processes is what you intuitively think it should be. But if you can actually make predictions based on that understanding, design new things, and have them work, that certainly gives you reason to be more confident.

Just to say it one more time: at each step, we have this memory module, and we want to make sure that, given a new token in the future, we can pull out the payload that we would have gotten if we were doing a full attention process. We do that with an approximation facilitated by this MLP, which allows us to return, given the key of the current token, the value of the current token. That's what these tokens have that leads to a certain payload.

Now, at runtime, I can take the query—what the new token is looking for. Those are definitionally similar, right? The key of the previous tokens and the query of the new token. It's that similarity that ensures that you get the right value from the retrieval process and can carry on from there with, approximately—in a fuzzy sort of way, in a more human-memory sort of way—not exactly what you would have had with full attention, but at least some close approximation of it.

That is, again, super, super interesting. For folks who want to develop their intuition about how these things all work and what's really going on inside, this has been one of the best things to take the time to understand for me in quite a while.

Let's talk about some details of that. It sounds like these things aren't super big at this point. What do we know about how much information you can compress into an MLP, and how do you think about how big this memory module should be versus how long you want episodes to be able to run, versus how long the sliding attention window should be? So far, how have you been thinking about how to size these things relative to each other?

Honestly, I think that the number of tokens we're using for the segment-attention part—or, for example, the number of tokens we're using for the memory—aren't really that challenging to tune. The way we want to tune these hyperparameters is not very challenging, because that's directly about how much we want to spend on these things. For example, do we want to train a model with hundreds of millions of parameters, or do we want to train a model with billions of parameters?

One important thing here is that there are some equivalent versions that we can use to understand how we can set these parameters for our design. Let's say we have a Memory as Context (MAC) or Memory as Gate (MAG) architecture. In that case, the memory part is running in parallel with the attention in some sense. It's very similar to a head: when we have a multi-head design, we can see that branch as a different head from the attention part.

One thing we can do is use the resources we have—how many parameters we want to use and all these things. When we know the answers to all these questions, we can use half the number of heads for the memory part and the other half for the attention part. That's one way we can do it.

For example, in the Memory as Layer (MAL) architecture, we can simply use whatever configuration we're using, but with respect to the number of layers, one thing we can do is use half of the layers as memory and the other half as the transformer, or something like that. As you mentioned, it's not very clear what the performance of the model would be if we changed these things. There might be a point where, for example, using 3 memory heads and 7 attention heads works much better than 5 memory heads and 5 attention heads, and all these things.

So basically, I think there is a lot of space to explore all these combinations and see which one is better. But generally, if you're talking about just using Titans instead of whatever model they're using that's based on attention, the configuration would be very similar. We can use half of the heads as the memory and the other half as the attention, or, if they're using something like memory as a layer, we can simply use half of the layers as memory: one memory layer, one attention layer, one memory layer, one attention layer, and so on. So that's another way that you can do it.

Do you want to take a little more time and describe the 3 different ways that this works? I was definitely going to make sure we talk about the momentum piece, or the surprise/momentum component to the update.

Yeah. Let's talk about the momentum part. I think that's the part before we talk about how we want to incorporate the memory into the architecture and combine that with attention.

Regarding the momentum part, I think we tried to explain the intuition in the paper. Basically, the intuition is that everything surprising is probably worth memorizing. But one thing we need to consider is that there might be a specific token that is very surprising to us, and the consecutive tokens are describing what was in that surprising token. Basically, all of them are important, so we cannot just ignore that one. But those tokens might not be surprising to us.

For example, let's say that I'm saying something like, “I'm leaving,” and this sentence can have a lot of meanings depending on the context. Without hearing the other parts, I cannot say whether this is surprising or not. Maybe I'm talking to a colleague, and it's the end of the shift, and I'm saying, “Bye, I'm leaving.” It's not surprising at all. I really need to understand what's going on around that specific sentence, around that specific token, and so on and so forth.

But, for example, assume that I want to leave the company. That might be surprising when I'm saying that I'm leaving. Depending on the context, there are some situations in which it is surprising and some others in which it is not. Basically, that's one way to think about it.

Another thing we need to consider is that other parts of the context that we need to pay attention to are somehow important, because those tokens might not be surprising; they're just simple descriptions of what we're saying about that specific surprising token.

How can we model this process? How can we say that this momentum part is very effective not only for modeling the momentary surprise of the token, but also for understanding how surprising the past tokens are? That's the part where we introduce the momentum.

The way we can think about that is to break the surprise metric into parts. The surprise metric is just a metric that says what information should be remembered and what information should be ignored. One simple way to define this surprise metric is based on gradient descent. But another way to make it more powerful is to break it into 2 parts.

The first part is the momentary surprise: how surprising is this exact token? This is the momentary part. There's another part that describes how surprising other tokens are. That's the past part, and actually, that's decaying. Let's say that there's a surprising token 64 tokens ago. We need to forget about that token over time, and that's different from the forgetting part that we have in the memory module. That's the forgetting part for the surprise metric.

A surprise metric also needs to decay over time, and the reason is that the context might change. For example, when time passes, these tokens might not be as relevant. Basically, we also need to have this decay mechanism. Again, that's one way to motivate these kinds of approaches.

If you want to discuss it technically, focusing on the mathematical formulation of that part, one thing we can say is that we have this design of test-time training. Now, instead of gradient descent, we can use a more powerful optimizer. For example, we can use gradient descent with momentum. Basically, these more powerful optimizers can lead to a more powerful architecture.

Another way one can say it is, let's use Adam. We use the Adam optimizer as the way to do that, and that would lead to a new architecture and potentially one more powerful than what we have, for example. So that's another way to think about this specific formulation.

One thing that I found is that there's a tradeoff between whether you want to focus more on the math side or whether you want to motivate everything in the sense that everything is compatible and all these things. Basically, I think the momentum part was a good spot to focus on to say that everything makes sense in terms of how humans work and how we can mathematically get better results. But definitely, yes, you can use other optimizers to make it work. Again, that's a different perspective on this matter.

Okay, so let me try to summarize this back, and tell me if I get anything wrong. Again, we have this memory module, the purpose of which is to allow, at runtime, a new query—a new token—to come in, take the query vector from that token, put it into the module, and output the relevant payloads that we had from earlier tokens. That is finite in size, so it can't just be more and more information in there forever. How do we update and manage this? Looking at equations 13 and 14 in the paper, it's actually relatively simple: we first decay the current memory state by a bit, and then we add an update term. That update term is determined by the loss from the current token plus a momentum term from the previous token.

The loss of the current token constitutes the surprise. The bigger that loss is, the more we realize we need to update, right? Because we did not predict that token well, we realize, “Geez, we really need to update significantly to be able to do a better job on this particular thing.” Then the momentum term says we need to maintain a significant update for some tokens to come.

We want to not just update on this one token, but this highly surprising token signals the beginning of a partial episode where we want to make sure we gather the information from that entire upcoming sequence. We want to make sure we pay extra attention to all of that—not in the technical sense of attention, but in the weight of updating the memory states. Intuitively, that feels quite right.

It is interesting that it also looks so similar to other optimization algorithms generally. Just on an introspective basis, it does feel like I can think back on moments in life where I've been very surprised. Then there's this modified mental state where I'm a little dizzy, off-center, but you often come away from those moments with very clear memories of what happened. I remember where I was when a certain thing happened, and I remember that morning.

The classic one, of course, at least for people my age, is that I was in high school on 9/11. I remember who told me that it happened. I remember where I was standing. I remember the class I went to next. I remember what we talked about in the next hour. So it wasn't just that one token, but it did create a similar memory.

I remember when I talked to my now-wife, then-girlfriend, later that day. I remember my dad came home from work. That whole episode, that whole day, is way more salient in memory than the day before and a few days after. It seems like you're essentially capturing a similar process here, or creating a similar process, with these update rules.

One question I did have is about the parameters. There's the rate at which the current memory state decays, and then there are the weights, the strength of the update. There's a free parameter—you have the loss, and then you have a free parameter that you multiply by that to determine how much to update—and then there's also another free parameter on the momentum term. Are those all learned but fixed? Am I understanding that right, or did you just pick them?

Just to make sure that I understand it correctly, do you mean the parameters like α_t and, for example, θ and α_new?

Actually, they are learnable parameters, and they're input-dependent. Basically, we project the input, and the model is learning how to project the input to parameters like α or θ. Based on the token, we decide which part of the past information is important. Based on the token, we decide whether we want to use the surprise from the past or not. Again, based on the token, we decide whether we want to consider this specific momentary surprise or not. All of them are input-dependent.

Gotcha. I should have gotten that from a subtlety in the notation, but I'm not always as great with notation as I'd like to be. So there really is a lot going on, and this is kind of a Mamba callback, too. It's not even something as basic as the decay of the previous state of the memory being dependent on the input at that particular time step.

So you really have a lot of ways for the input to influence how the ultimate state of memory is being updated at each time step. It's a super-expressive, I guess you might say, setup in that sense.

Cool. I think that is again really interesting, hopefully pretty intuitive. How hard is it to make all this stuff work in a scalable, efficient way on the given hardware? You are doing something that hasn't been done much, where you're doing gradient descent at each time step and making these updates. I don't really have a deep sense of how hard that would be to manage. I don't know if you're using GPUs or TPUs, and I'm also not entirely sure how much difference that would make in terms of how hard it might be.

How much low-level wrangling, manipulating all the relevant vectors and matrices, did it take to make this scalable, and how scalable did you ultimately achieve in this portion of the work?

Yes. Actually, we have some experiments in the paper. One thing that we can see is that if we implement these kinds of approaches—updating the gradient descent at each step, and so on—then the model can become a recurrent model. This recurrence can be very slow in practice when we want to train the model. But the good thing is that this process of calculating the gradient and updating the model can be reformulated in matrix-multiplication notation. Basically, this matrix multiplication can be very fast using GPUs and CPUs, so the good part is that we can make this process faster and parallelizable.

The TTT paper discussed the dual form of how we can, for example, write this recurrent formula for gradient descent in matrix-multiplication format. But we have additional modules that we need to consider. The first one is the weight-decay part. When we add weight decay, we need to multiply the memory by a constant number, by a scalar, or, for example, by a vector if you want to do channel-wise decay. The interesting part is that even using this decay formulation, we can again have a matrix-multiplication format, but we need to add additional Hadamard multiplication or, for example, construct a diagonal matrix and do the matrix multiplication for that one.

That's generally the way that we can make the process parallelizable when we use the decay part. But we also have the momentum part, which is much more challenging if we want to handle that. The interesting part is that, if we look at the formulation in the paper about momentum, as you mentioned, the momentum is something like S_t, the surprise matrix, equal to the decay parameter times the previous state minus the gradient part. Interestingly, this is again a recurrent neural network—a linear recurrent neural network.

There are different ways that we can make this process parallelizable. The first one is using a scan algorithm, a parallel scan algorithm, because this formulation can be very similar to, for example, Mamba. It's a linear recurrent model, and we can use a parallel scan to calculate all the surprise metrics in one chunk. That's one way that we can do.

Another way is to use the matrix-multiplication format for this linear recurrence as well. Let's just expand the surprise metric part, the recurrence of the surprise matrix. Again, we can see that the gradient part can simply be reformulated as matrix multiplication, and also the coefficient of theta can again be formulated as diagonal matrix multiplication. That's another formulation that we can have to make the process faster.

There are different ways that we can incorporate these techniques into the training part. Actually, these techniques are very effective because, in Figure 9, we compared the training time of Titans, some modern recurrent neural networks, and attention. We can see that when we increase the context length, attention will definitely drop. That's one motivation for using these linear models.

Compared to other models, we can see that, for example, Titans is faster than Mamba, but there are some other modern linear models that are somehow faster. But one thing that we need to somehow emphasize here is that, in the Titans paper, we focus on delivering a new perspective and a new architecture, but we didn't spend so much time optimizing the implementation—for example, using some kernels to make the process very fast and all these things.

Basically, our focus was on the architecture side: designing new memory modules and also delivering a new perspective in that sense. Definitely, in the future, it would be very interesting to see how we can design different kernels to make the training process faster. In that case, I really believe that we can achieve comparable efficiency even to simple linear RNNs that are very fast. So, yeah, that's generally about the efficiency part.

Cool. Yeah, it's off to a great start, I would say. No doubt there's still room for further optimization. Let's talk for a minute, I guess, about the 3 strategies for integration. You've kind of alluded to this a little bit: there's memory as context, memory as gate, and memory as layer.

I think memory as context and memory as layer are probably the 2 most intuitive. Context basically being: do the retrieval, get the information out of the memory module, put that into context, and do essentially normal attention. We've seen many things like this, even including multimodal, right? If you have a fusion of a vision model and a language model into a vision-language model, sometimes there's separate encoding of the image into its space, and then the text can go into its space, and later they can have attention that makes sense of them together. That's kind of how I've read memory as context.

The memory as layer is also pretty interesting. It's a little less intuitive, in the sense that I feel like, especially as you were talking about interweaving the layers, it's a little odd to think about accessing memory multiple times through the forward pass, and especially as you get to the very late stages of the forward pass, still going back to the memory. It's not super intuitive to me how that would work or why it would work that way. And then there's also memory as gate, and I probably have the least intuition for that. With that prompt, can you help me develop intuition for all 3?

So let's say that we have this long-term memory, and we are using attention as the short-term memory. Let's just ignore the persistent one, because the way that we are treating the persistent memories is always the same. So let's just ignore that part. We have this short-term memory, and we also have this long-term memory. As I mentioned, it's not very clear how we can connect this short-term and long-term memory, even from the neuroscience perspective. There are different ideas in the literature.

For example, memory as a layer is the most common way that people do these hybrid models of RNNs and attention. What they are doing is simply using each of the memory or RNN and attention as one layer, and so on and so forth. I think one of the important messages that we have in the paper is that most of the literature is focusing on one configuration, which is not great. We have 2 different configurations, or 2 different types of architecture, that we can use, and they are actually more powerful than the architecture people are using in the literature.

That's an important message here, and actually there might be some configurations that outperform the other configurations. Generally, one thing that I can say is that these kinds of approaches are worth exploring. That's, I think, an important message here.

But what is the intuition behind each of these designs? Let's say we have something like short-term and long-term memory. The data comes in. One way to model this process is to say that our short-term and long-term memories are interconnected. So what does that mean? It means that short-term memory will help the long-term memory, and long-term memory will help the short-term memory.

How can they do that? The data comes in and goes to our long-term memory, which extracts the relevant information from the long-term memory and then uses that as the context. When we use that as the context of the memory, one thing that we have here is that the attention part will decide whether it wants to focus on the current context or whether it wants to focus on the past information from the long-term memory.

The output would be something that is somehow combined or compressed by the attention. It's a piece of information that goes through the attention, and basically we know what information is important and what information is not important. So the output of attention goes to the long-term memory and lets it know what information should be stored in the long-term memory.

So that's how we can describe this MAC architecture. As I mentioned at the beginning of the process, the long-term memory is helping the short-term memory by extracting the past information that can help the short-term memory decide how to combine the information or, for example, learn from the data. On the other hand, the output of the short-term memory will help the long-term memory understand what information should be stored and what information should not be stored.

And so that's the MAC architecture. Another way we can think about short-term and long-term memory is to treat them as 2 different modules, 2 different branches of memory. Basically, in this design, the data comes in and then goes to long-term and short-term memory at the same time. We have the information from the long-term memory that is related to this specific input, and we also have the information that comes from the short-term memory.

How do we understand the pairwise interactions of all these things in the context? We then just concatenate these 2 or multiply these 2 at the end. We are using this long-term information and short-term information, and then just concatenating or multiplying them to use both types of memory that we have. That's basically the main intuition for how we want to do these things.

Finally, there is memory as a layer. The main reason we put memory in a layer is to have a unified way to combine the memory and attention parts, because most people in research are using this configuration, and it can help us understand which types of configurations are more helpful.

The way we can think about that is, let's say that the data comes in, and what we are saying is that short-term memory and long-term memory are basically modules that are sequentially connected. The data comes in and goes to, for example, our short-term memory or the previous state of the long-term memory, depending on how we want to order these layers.

For example, the data goes to your short-term memory, and the short-term memory will decide which information should go to the long-term memory. We repeat this process. The output of the long-term memory can again help the next layer of the short-term memory, and so on and so forth. In this case, each layer of our short-term and long-term memory is helping the next layer of long-term and short-term memory. That's the main intuition behind thinking about memory as a layer.

Yeah, it is interesting to observe, just looking at all the results in the paper, that, as you said, for the 3 hybrid approaches, memory as a layer is winning very few categories. That, I guess, again speaks to the motivation—or the higher-level thinking—being an important way to approach this, because memory as a layer is very similar to a lot of the hybrids we've seen between attention and Mamba or state-space models in the past.

I think the way that came about was that the original Mamba was just all Mamba—there was no attention in it at all—and then people said, “Geez, what if we can get the best of both worlds by just interleaving these things together?” That also seemed to improve on the base, but it wasn't necessarily a super-principled thing. It was just, “I see that I can stack a bunch of attention layers, and I can also stack a bunch of Mamba layers. Now I can shuffle them in together,” and, lo and behold, it gives me, in some sense, the best of both worlds.

But it's interesting to hear that memory as context and memory as gate both have a more principled, higher-order rationale motivating them, and indeed they do outperform the layer approach almost across the board, though maybe not entirely. I'm just roughly counting, but it looks like 9 different categories and 3 different scales at which these experiments were run. Memory as a layer wins maybe 2 of roughly 30 of those different categories, and then between context and gate, it's maybe 50/50 across all the rest.

One big question I have, and one of the things that got me most interested in the Mamba architecture and hybrids in general, is that it seems like there are different—what I've started to call—micro-skills across these different architectures. Transformers, for example, are much better than Mamba when it comes to retrieving or repeating past patterns. They can see the past, which is not intuitive, because they can see the past pattern explicitly, and so they can repeat that past pattern as needed.

Mamba has some other micro-skills that attention struggles with, including learning really sparse signals or signals in really noisy environments. Sometimes transformers can struggle to learn those, but the Mamba architecture seems to do better.

Do you see any micro-skills? I don't know if you've had any ability to study this, but are there things that this new mechanism can do qualitatively differently or better, or perhaps things that it can't do as well? I guess we'd probably have a pretty good sense of that, since it doesn't have the full thing in explicit context anymore. But have you seen any micro-skills that are new, different, exciting, or even just informative?

I think one thing we need to consider, if we want to answer that question, is whether we're talking about the neural memory module part of Titans or the entire architecture of Titans. If we just focus on the neural memory module, one thing we need to consider is that it's similar to linear models. It's very similar to RNNs. It probably has the properties of what we know about RNNs and these kinds of approaches, but it's generally more expressive than other RNNs that we have, because the memory is more expressive, the architecture that we are using for the memory is more expressive, and the recurrent formula is more expressive. That's the main advantage of that neural memory.

The entire architecture of Titans—for example, the MAC architecture or MAG architecture—consists of hybrid approaches. They somehow have the ability to have the best of both worlds. They can use the advantages of transformers and RNNs, at least on paper, because we can ignore the output of the attention when the task is RNN-specific, or vice versa.

That's one thing we need to consider. On paper, as I mentioned, theoretically, these hybrid approaches can have the best of both worlds, but there are some cases where we cannot get better results than a pure RNN architecture. Similarly, for attention, there might be some cases where we cannot achieve better results than a pure transformer architecture. The question here is whether there are any specific tasks where hybrid approaches are better than transformers and RNNs.

We have another paper, and it's actually called The Best of Both Worlds. We see this process through the lens of graph algorithms. We want to see whether, for example, a hybrid approach can have better results on some specific tasks. It turns out that there are some specific tasks where using a hybrid approach is better than using a pure RNN or a pure transformer-based model.

One thing we can say is that Titans-style architectures can be connected to RNN models, and theoretically, they are more powerful and expressive than existing models. For example, models like Mamba, RetNet, S4, S5, and all these models are limited to the TC⁰ class of problems. There is a paper called The Illusion of State in State-Space Models, I think. Basically, they have shown that state-space models—or, more accurately, diagonal state-space models—and transformers are limited to the TC⁰ class of problems.

On the other hand, if we go beyond diagonal to linear RNNs with nondiagonal transition matrices, similar to DeltaNet, we can see a more expressive architecture and go beyond the TC⁰ class of problems. To say more about what a TC⁰ problem is, one example is the state-tracking problem. Let's say that we are at the 0 point, and there is a string of actions: go left, go right, and something like that. At the end, the question is, “Where are you right now?” The model is expected to understand these things.

Diagonal RNNs, including state-space models, and transformers are limited and cannot do that task, but a nonlinear RNN can do it with only 1 layer. It's a very simple task for them, but linear RNNs and transformers cannot do that. Some models, like Titans—the neural memory without the full architecture—are also on the nonlinear RNN side, and they can do this state-tracking problem. Basically, in these kinds of tasks, they are more expressive and powerful than the other approaches that we know.

There are some other RNNs that are also capable of state tracking. For example, DeltaNet is capable of doing that because it's not diagonal. But models like Mamba, Mamba-2, RetNet, and S4 are limited, and they cannot do this state-tracking problem. These nonlinear RNN approaches, similar to what we have in Titans, can do that, and linear RNNs with nondiagonal transition matrices, similar to DeltaNet, can do that as well.

Is that analysis limited to a single forward pass? Because I would assume that I could give o1 or similar reasoning models a problem like that, and I would be pretty surprised if they couldn't talk their way through it.

One thing we need to clarify is that some of these approaches might work well in practice. When we want to understand them from a theoretical perspective, it's really hard to model everything that we have in the real world. Basically, we need to make some assumptions that make the process simple.

So we could use this theoretical understanding and these theoretical frameworks to develop some results. Yes, I think in practice large models are definitely capable of doing some of the tasks that I mentioned. But when we're talking about theoretical results, it's definitely not completely about how the models work in practice. Actually, the paper that I mentioned, “The Illusion of State in State-Space Models,” also has some experimental results that support the claim.

I think that's the main thing that we need to consider. There are definitely some other combinations of models or, for example, additional techniques—reasoning models and all these things—to make the model more powerful and able to do a lot of stuff. But in these theoretical frameworks, we're just talking about the simple model that we have.

So basically, that's about the RNN side—the recurrent neural network side—which is our neural memory. The attention part is also well studied in the literature, so we don't have any specific contribution on that side. Attention is attention. As I mentioned in the best-of-both-worlds paper, we also have some results showing that hybrid models are better at doing some tasks that are hard for both RNNs and attention, and that their combination can be more effective. That can also be another motivation to consider these kinds of hybrid approaches.

Cool, yeah, I'll need to go check that one out a little bit more deeply. It sounds like, if I'm understanding correctly, there's a distinction in some of these theoretical contexts, and maybe this is a good thing for people to be more aware of in general. If I'm understanding it right, it sounds like some of these papers that say transformer architecture can't do X are perhaps better understood as saying that it can't provably and reliably do X under certain conditions all the time, but maybe still, in practice, can do it much of the time.

Maybe without a guarantee, but with enough inference tokens to burn, you don't necessarily need a guarantee. You just need consistent success, even if it's not fully theoretically proven that you will always get that success. Those different frames may be why there's a decent amount of people talking past each other online about what exactly can and can't happen.

It is often strange when you see these things where it's like, “We've proven that a transformer can't do X,” and then somebody goes to ChatGPT and does it, and you're like, “Wait a second. These two things can't both be right.” It seems like it's probably often a question of different assumptions and different levels of what they are considering to be proof. The robust-guarantee level versus the “most of the time it works” level can be quite different.

I guess, just to follow up real quick on the micro-skills concept, when you look at all the different things that you tested the new Titans architectures on, what jumps out to you as the most exciting thing that it is doing better than other architectures?

I think the most exciting part is the long context. Actually, I need to say that the long-context part uses some synthetic datasets, so we might not be as effective as what we are showing in the paper when we're talking about general tasks and a general model. Basically, what we are doing in the paper is comparing the model with some counterpart approaches and seeing how they're performing on specific tasks so we can compare them.

The interesting part, I think, is that Titans can outperform other models in long-context tasks. This is very impressive, in my opinion, when we can see that they can scale to 2 million tokens. That's something that some models like GPT-4 cannot do. The performance on the benchmark for GPT-4 can drop very quickly, but Titans, with a very small number of parameters, can scale to millions of tokens and even go to 10 million tokens with accuracy of about 70%.

I think that's the most impressive part, in my opinion. For long-context modeling, the MAC architecture is very important. It's very important to have 2 different branches for long-term and short-term memory, and these 2 help each other understand the context. When we use other approaches like MAG or, for example, memory as a layer, we have much shorter context lengths with higher accuracy. Basically, the MAC architecture is very important for long-context modeling.

The behavior of all these models can be different when we use different-sized models. For example, there is a very great plot in the paper called “Mixture of Mamba,” I think, if I recall correctly. It's a plot of the perplexity of Transformers, Mamba, and a mixture of Mamba. One thing we can see is that the RNN approach—in this case, Mamba—scales better when we have less data, but when we increase the number of tokens that we use for training, Transformers scale better and outperform the other approaches when we have so much data.

That's a very important thing to consider. One thing we can see about Titans is that, at least in what we experimented with, we couldn't see that pattern. They consistently have good results compared to Transformers and other linear RNNs, which is definitely a good thing for converging faster.

But when we're talking about using additional data to train the model, we might see some ups and downs in the ordering of models with respect to perplexity. I think that's also very important, which might not be very clear from academic papers, because we need to have a fixed number of tokens that we're using to train our model. It's generally infeasible to train the model on larger and larger numbers of tokens.

I think that's another good result that we can see in Titans: They consistently outperform the other models when we increase the number of tokens that we're using in training.

So, as we look ahead, how much do you think data is going to become the limiting factor on long-term memory? When you show that still a very small number of the architectures can even attempt to get out to a 10-million-token context, and of those that can, Titans is blowing away the other 2 that you show in the experiments, it strikes me that we don't really have many sources of data that are actually such long episodes.

The internet is full of short, medium, and even long blog posts, but they're not reaching into the millions of tokens very often at all. If we think about wanting to have long-running agents, it seems like an architecture like this could be really key to getting a lot of the behavior that we might want out of a longer-running agent. But we still don't necessarily have long-horizon datasets to train on.

What's your take on whether we have the data that we need, or whether we have to go create it if we want to really take full advantage of these sorts of architectures at scale?

I think this is related to 2 concepts in the community that are usually important for RNNs: length generalization and length extrapolation. The question is whether our model can somehow work better, or at least not show any performance drop, when we increase the sequence lengths.

As you mentioned, most of the data that we have might not be very long-context data. It might be challenging to train the model on millions of tokens, but we actually don't need to do that. Generalizability in some models can be important for effective training. Let's say we train our model on a dataset that might not have a very long context, but we expect our model to understand how to generalize to a longer context.

When we have a single model like an RNN or, for example, our neural memory, we can see that these models struggle with length generalization. When we increase the sequence length, we can actually see the performance drop, and that happens very soon—at 16,000 tokens or something like that. It's not at 1 million tokens.

With Transformers, we can see some similar results without doing any tricks to make them better at length generalization. But there are different architectures that can somehow mitigate these issues and make the model more generalizable with respect to sequence lengths. There are also some other tricks that we can use—for example, different positional encodings or different normalization techniques that people use for length generalization.

Architecturally, I mean, regarding the architecture design, we can see that the MAC architecture is very good for length generalization, based on my personal experience with it. My understanding is that when we have something like the MAC architecture, the memory module is only responsible for summarizing the data for the attention into a fixed-size number of tokens.

So it doesn't have a very hard time compressing the information and learning things. On the other hand, attention doesn't need to attend to so many tokens. We can just use a fixed-size local attention. So, again, it doesn't have a very hard time learning from that part and understanding how it can use the information from the long-term memory.

In general, one thing I can say is that having more data is definitely very good, and it can help us have a better model. But it's not the end of the story. We need to have better architectures and better techniques that we can use in training the model, so we could make them more generalizable—for example, make them more generalizable to different sequence lengths and make them better at longer sequences, and so on. I think architecture and the model itself are very important.

I guess one question I have is: how retrofittable do you think today's open-source models are with this technique? If I wanted to take an off-the-shelf Llama model or R1, or what have you, and integrate a long-term memory module of the sort that you have, it feels like I could probably do that without—I don't want to make it sound simple—but based on what you just said, it seems like I could probably hack the Llama architecture a bit, bring in the memory module as you've designed it, and then do some amount of continued training to get the memory—to get the thing to actually effectively pay attention to the memory as context. I kind of think that that will work. Are you expecting that to happen?

I think potentially it's possible. The main thing that I like about the MAC architecture is that, no matter how large you want to choose your context lengths for training your attention, you can always go farther with the long-term memory part. You can use attention with context lengths of 32. You can use 128, 512, and go beyond that, like 8,192 context lengths, and so on.

No matter how large you want to choose the context lengths for training the transformer part, you can always go farther by using additional memory. So, with the memory-as-context design—the Titans-style model—I really hope that these kinds of approaches can be used in the future and basically become something that we can use, because they are actually working and can help with performance in long contexts without, for example, a performance drop or damaging the performance of pure attention. I really think that this kind of future work is very promising.

One other conceptual question I had goes back to how the memory module is predicting the keys. It struck me that one of the promises of long-term memory, at least in humans, is that we can learn whole new domains of stuff that we've never encountered before. If we spend time on it, it can become part of our background world knowledge.

Here, it seems like we still have a memory module that outputs something that is, by definition, in the latent space of the model as it was originally pretrained. Because those main weights don't get updated at runtime, the memory module is always outputting something that's in the space of that pretrained model. That got me thinking: would this ultimately be something that people might want to combine with continued pretraining?

If, for example, you're an enterprise—let's say I'm 3M or GE. I'm some company with a long history and millions of products, and just some super-deep history that I want my models to learn—it seems like I might want to do continued pretraining so that the model itself has more representations of the things that I care about. Then that might be needed to get that long-term memory to be able to output things in that modified space.

Does that disconnect between the model's fixed latent space and the fact that the memory module is outputting into that space suggest a frontier for future elaboration? Perhaps that continued-pretraining approach would solve it, but I'm interested in your thoughts on just how far you could push this. Could you get the model itself to learn whole new domains of knowledge this way, or would you expect that you'd bump into some limits and have to do something to overcome that?

I think that if we want to train the model for a long period of time—and by training, I mean even test-time training, updating the parameters of the model over time and using that for different tasks or different types of things—we can face a problem that is very well studied in the literature called catastrophic forgetting.

Basically, one thing that we can see is that the model learns a specific task or, for example, a specific context. When we want to go to the next task and learn it, the model might completely forget the previous task to get adapted to the current configuration and task that we have.

I think the part you mentioned is very promising for future study, but there are some challenges, and catastrophic forgetting is one of them. If you want to have a very long time in the test-time training part, the model might forget about the initial tokens or, for example, the previous task that it was trained on, and all these things. That's really challenging.

If we could solve that, we would solve all the robotics stuff and all the reinforcement learning stuff and all these things. That's a very challenging question, but I think we might face this issue, and it needs to be addressed.

I know we're just about out of time. Any closing thoughts about where you are most excited to see all this stuff go next? To me, this feels like a pretty notable step on the path to really figuring out that kind of continuity of memory. I do think that's going to unlock some of the most exciting applications that people are interested in.

Whether it's the drop-in knowledge worker that has that full context of your enterprise's full history and everything that it's ever done, or the long-running agents that can go out and make mistakes and maybe learn from those mistakes and not make the same mistakes twice. I don't know that this gets us all the way there, but it feels like it does take us a meaningful step forward. So what, to you, are the most exciting or promising directions that you hope you or others will go in from here?

Honestly, one thing that I really think is interesting is using these kinds of approaches for other tasks and modalities. We're trying to do these long contexts for text, but sometimes there are other modalities or other tasks that are more suitable for long context. For example, as you mentioned, agents or reinforcement learning, decision-making—all these things require some long-term memory to understand different types of patterns in the data.

I think it's really promising to see how this long-term-memory-style architecture, like Titans, can work in those domains, as I mentioned: decision-making, reinforcement learning, and other modalities and data modalities. I think that's a very promising future direction—to see whether these designs that we just talked about are actually effective outside of language modeling, and whether they are just good architectures for language modeling or good architectures in all areas of deep learning.

That's the important part, because I think one of the main reasons that Transformers were really successful is that they were successful in different domains. I think that's really important to explore and see whether these additional modules that we are adding, and new architectures like Titans, can also be effective in other domains, other data modalities, and all these things.

Definitely, there are so many ways to modify them for different tasks. For example, the design of the MAC architecture might be very suitable for language modeling, but the MAG architecture, using gated short-term and long-term memory, might be very effective, for example, for decision-making. We don't know that. All of them are really interesting for future work—to see which one works better and which one is effective.

Well, it never ends. There's always plenty more to do, at least until AGI takes over and starts doing it all for us. But for now, this is a really fascinating development. I think it's a great combination of high-level, principled, intuitive motivation for the work and obviously some really deep work in terms of making it work on the mathematical and even computational level. It's a pretty impressive piece of work. I've really enjoyed studying it, and I'll certainly be looking forward to what you guys come up with next. For now, I will just say, Ali Behrouz, thank you for being part of The Cognitive Revolution.

Thank you very much for having me.

Thank you.