Titans: Learning to Memorize at Test Time (Paper Analysis)

[snorts] Hello. Today we're going to look at Titans learning to memorize at test time. This is a paper by Google research and has been pushed as part of their Nurips publications. So there's a blog post about it and people posted on social media about it and um I have to admit I just fell for the marketing. So here we are. I got interested in this paper and it is a good paper. uh it so it proposes a architecture of a um model u so think of like let's say a language model right that learns to memorize at test time so it can go beyond the current context you can go over a very long piece of text for example you can slice that into parts and you can go over these use this memory to connect to remember stuff from earlier parts and carry them to the next part thereby overcoming the context window limitations uh that are typically associated with something like transformers. So that is very cool. However, I do find that sometimes they um kind of reinvent old things and um present them as maybe a bit more novel and then other times they just assign names to things to make them seem like they're a new thing. So, a lot of the things they call memory in this paper are basically things that already existed. Um, but we'll get to that. So, I'll say it's a 50/50 split of cool new stuff and, you know, maybe putting a bit the foot on the on that marketing pedal. So, let's dive in. the we're the paper is um also of a decent size and we're just going to tackle the what I feel are the important bits. Um but if you have a different view then um that's your right leave a comment and tell me about it. So the problem I've already outlined here in the most models nowadays like either they are natural sequence model like an RNN an LSTM or something like this but these models have been largely surpassed by attentionbased models like transformers and uh the problem is that the transformer can only attend to things that's in the current context window um and that context window can become very large for some tasks. So think of things like video understanding or you know very long real world tasks right where a lot of things happen and you basically have to consider all of these things in order to make your next move. So the problem is quite simple. you have some long piece of data and your model only has the capacity to look at maybe data that is you know this oopsie that is this long here. So let's say this is the size of the context window of your model. Now you can place that anywhere like you could place it here right and so on. Um, but you cannot push everything into one context window. That would just blow your model out. Um, and so a lot of people have been trying to tackle this like back to the earliest variants um of that followed BERT. People have been trying to come up with u variants that would go to very long context. Um, some of the variants back then make explicit use of something that you might call nowadays memory. So where they say okay let's first consider this part right here and then at the end let's produce some artifact like something and usually that thing is just kind of the the key like the some kind of computation um during the last token here because naturally attends to the whole context right so it kind of integrates something here so let's just take that hidden state and kind of make it available to the next context window so that next context window. Um, when it produces this token here, it cannot attend. So, if it wants to attend to something here, it cannot do that. But it can attend to this thing here. And this thing here supposedly is kind of like a compressed variant of the whole context of here. And now obviously if this thing again has from the last chunk one such thing, then you kind of get a get this also in this. So it kind of acts like an RNN in that it has this kind of state that it passes on to the next thing. Whereas within the context window it acts like a transformer. These were like I don't know things like Xformer or something like a transformer X maybe or transformer XL. I don't

remember the exact names but a lot of ideas like this were around. um other ideas and that's mainly the ones um the ones discussed here are ideas around linear transformers. So linear transformers uh there's obviously the basic variant where you just say well instead of doing all that softmax shenanigans let's just make linear operations but then there's also the bit more advanced things that you would find in things like I think like performer or reformer one of them um was very prominent in that where basically you using kernel functions uh in the attention mechanism in order to um pull it apart. So the way this works is you have your attention and I think they have down here somewhere, right? They have the formula for attention. Now um basically what is what is this? Okay, you have your tokens and let's say you produce this next token right here. So you would say okay, I want to produce this token. let me um create the query from this token and let me multiply it by all the keys from these tokens right here. So you produce these queries and the keys and the query is just going to attend to all the keys meaning it's going to create the inner product of the query and the keys and um so the inner product of this query with each of the keys right that gives you basically how much the query matches with each of the keys and then you're going to normalize that using a softmax operation. So, um, what you're going to end up with is kind of like this distribution, right? So, okay, like key 1 2 3 4 5 key one it matches this much, two this much, three this much and so on. Um, after the softmax, this all sums the sum of this is one, right? This is normalized. So now what you can do is you can use this as basically a waiting function for your values and the values are also computed from this right. Um so you can consider the attention mechanism kind of like this here actually gives you the weights that you use to aggregate the values and um so it's kind of like a um fully connected layer except the weights of the fully connected layer aren't fixed. They are data dependent. and they're sort of computed on the fly by this query key interaction. Um so the problem here is actually this softmax uh this softmax is an operation that is nonlinear and you cannot pull it apart right you cannot um if you have a long context you still need to compute everything inside uh because that softmax operation um is not like separable in an easy way. So it's not scalable Now things like fast uh sorry flash attention they manage to kind of trade off um some speed for memory or vice versa but there's no like um closed form nice reformulation you can do uh this isn't the ca so ish uh because if you're familiar with uh kernelization and the whole kernel literature um you will notice that okay let's generalize this and let's say okay there's just some nonlinear function here it doesn't need to be softmax just some nonlinear function uh that applies across these things and if under certain conditions right um this nonlinear function um if it is a kernel this nonlinear function actually describes a distance between these things in this case it kind of does Right? I'm not sure how the normalization plays in here, but in this case, it kind of describes a distance between these two. And that means by the whole kernel literature, which I am not competent in, nor could I recite on the spot, but um if you are old school machine learner, you will have learned this at some point. Um then this kernel um well I believe actually this let's say this is like K instead of the soft max so K is just a general function um if that is a kernel that means you can actually like K of Q and K you can separate it into um separable function and this then becomes a true inner product between um 5 Q and five of k. Now what is phi obviously or s or whatever this theta whatever this letter is um what is this? This is a projection

a into some a nonlinear projection possibly into some space and in this space um computing this kernel is equivalent to doing this inner product. Now usually we are interested in this direction. Usually we're interested in saying we would like to achieve this inner product between things. Um but you know that's way too high dimensional right sometimes these are infinite dimensional spaces. So why don't we rather just compute this kernel function um because that's you know we can actually evaluate this. However, in the linear transformers case, we're interested in the other direction saying like, hey, look, this soft max here, if we just compute it out, right, this gives us this quadratic blow up and so on. If we could if we knew what this function was for the soft max, right? And the space with where we have to send it to make these inner products, we could just separately compute this of Q and this of K and then it's just an inner product. And then you know this whole quadratic blow up we could just separate it because it's a linear operation. And so that's where that's how these linear transformers attack this. They basically I think we have it here. Yep. They basically say okay if we could separate this right and we can do so um under the assumptions that uh this function here is a kernel then we could just pull this sum here that now goes across everything. Right? we have to first compute this and then do the sum and the normalization. Well, no. Now we can basically do the sum inside right here and only then multiply it by the Q and the normalization is also outside of this of the sum and so that becomes computationally feasible because effectively it's just a linear accumulation. You can see here the keys and values. That's your past, right? That's your that's the past. And you can just accumulate that. And if you go to the next token and you have the key and value, all you have to do is you only have to compute this thing here and then add it to the sum that you already have. And because you only need this sum here, you don't whereas with in your original attention, sure you can save the keys and the values, but you will still have to do all of this for all the past, right? You multiply the queries by every single one of the key. Then you do your nonlinear operation. Then only then can you multiply by the values. So that's the big thing here. You can multiply by the values first and just accumulate it into one sum and then you can multiply the sum by Q and that's ultimately what saves you all of that what avoids that quadratic blow up. Um only works because it's linear. So that's why they're called linear transformers. Um well okay I believe in history a lot of things were called linear transformers. Um uh Jurgen Schmid Huber is running around and saying that he invented transformers with an asterisk which is he invented unnormalized nonlinear transformers which just means inner products. Um and I'm pretty sure he did not invent inner products. However, this is the variant with linear transformers we uh worry about in this particular paper. Now what they're saying here is that okay actually linear transformers good in concept like good in principle however um I don't think the softmax has an actually computable function so these the things that we can do are just kind of approximations to the softmax function and um turns out these approximations are actually not good uh because linear transformers aren't very performant nowadays like they're not very competitive Um, and the paper here says, well, the linear transformers do not show competitive performance as the kernel trick makes the model a linear recurrent network. That's what we just discussed in which the data is compressed into a mat into a matrix valued state. Um so they're saying okay because all the data is compressed into one matrix value state is basically the same as an RNN where you have to compress the entire past into this one vector valued state and it's just impossible to kind of like a very long context cannot be compressed in a small vector valued or matrix valued state and because the

transformer doesn't compress it keeps the whole history and then it always computes um the interaction with each of the history. It does not have this problem. This is the first part where I'm skeptical honestly because like I think this is just a consequence, right? Like the linear transformers could also just keep the whole history and then compute all the interactions. It just turns out that it's formally the same as accumulating them. Right? There's nothing about accumulation that is the bad part here. It's about the the like you're choosing a function that is worse um to do the interaction with the past. Likewise, if we could like if we knew the kernel expansion function like the up projection for the softmax and we could do so accurately, we would absolutely do that. And then we would like compress like certainly there is a way to compress the entire history um into a matrix valued state even in a regular transformer. We just don't know how to do it. uh and then there is a way to interact with that matrix valued state in order to achieve some goal. We just don't know how to do it and um and it's going to be a nonlinear interaction right with that matrix valued state but if we knew how to do it we would absolutely do it. So to me it's not necessarily about the fact that it compresses anything that is the problem here. Um, but this paper really wants to make this point because they're saying like, well, we shouldn't use matrix valued or vector valued memories. We should use actually a neural network as the memory. And that to me is not a bad idea, but um, so we're going to jump a whole bunch of things here because I want to get to the point. What they propose here is they're saying well let's just generalize all the architectures um and let's just say actually all of these things they basically just like you can see all of these architectures as being sort of having a memory at runtime. For example, a transformer the memory is just the keys and values of the sequence so far. So the memory grows over time, right? In an RNN, the memory is just the hidden state that's, you know, carried from one uh iteration to the next. And that to me is kind of like, okay, you can describe it like this. And then they're saying, let's make that memory more powerful. And so what they come up with, and I think it's best um described in this picture, is here. The memory is actually a neural network, right? The memory is a neural network and uh let me just draw it. So you're you have your sequence here, right? Doo and you produce your next token. And every time you have this other thing here called the memory and every time when you compute this token here, right, you have your query and you can attend to the current oopsie to the current context, but you can also go to the memory and the memory will give you back stuff. Now I'm deliberately saying stuff here like but the memory will give you back information about the more distant past than your current context window. I mean maybe it can also give you stuff about the current context window but the idea is that the memory gives you back stuff about the more distant past. So now when you do your attention here you can um yeah you can attend currently but you can also um attend further back by going to the memory. Now if this memory is just as they say oh it's just a matrix valued memory right then what you would do is the way the memory would give you stuff is you just multiply the query by the memory and then that gives you the um or maybe also with the values. the memory is like the keys and the values. Um that would give you back your data here. Um this is actually exactly what happens in the linear transformer, right? This is the memory is just the sum of the keys and values. You multiply the Q and you'll get back the result of that. So this is basically what happens in the linear transformer. Um if you yeah if you go if you think of a real transformer and let's say you just have really long context the memory cannot be represented with a multiplication like this. So you need to do um something else and in

general they're saying well the memory should just be like a function where you plug Q in and it gives you back that data that you need from the past. So the memory is something you stick your query into and you get back data for your current inference. And now obviously the question is well how does that memory look like? And they say well it should be a neural network. network that kind of learns to at test time learns to kind of memorize and compress the data as you are processing it and keep it for a longer time than the current context window. And so let's say this token here goes out of context at some point. The memory, you know, as it was processing, it was learning from all of these tokens. the memory should still be able to tell you something about this token right here and that's the concept and so the memory in their case is like a two-layer neural network I think okay you have like your fully connect it's just an MLP these are fully connected layers right your Q goes in here and then your whatever value not referring to the values necessarily of transformers but whatever your resulting value um comes out the other end and So what does it mean to memorize in this paradigm? To memorize means to basically update the parameters of this neural network. Right? So we have an inner learning problem in this meaning at test time we kind of initialize for each let's say we're processing a really long sequence of text right at the beginning we would initialize this memory to some initialization um maybe random maybe fixed initialization and then as we are processing maybe we have a chunked um processing with a transformer as we are processing we're basically updating this memory meaning that we're training ing this neural network on the fly during test time so that later chunks can use it to retrieve stuff. And so the question is obviously yeah how do we do this? And they basically propose to train it with the following loss function. Um they train it to just rem to just remember the associations between the keys and the values. And so they train it because as you are computing these you know individual tokens you always get a key and a value right and what you want is you want basically you want that um if the query is similar to the key it's it should spit out the value. That's what the traditional attention does, right? If the query and the key are very similar inner productwise, it would spit out the value. And so it's quite natural to say, well, let's just train it at test time so that the if I stick in the key, I want to get out the value. So my loss function is basically that um the x is the key and the y is the value and I train it. And that means if I do this well, it means if I way later in the sequence have a query that is very similar to this particular key, it will give me this value right here. And so that's how I remember the past. Okay. Um and then they say well they arrive at this loss function um and they say well basically this is the loss function. So what we want to do is we want to train this neural network um as a function of its surprise. meaning that oh if you know if an event is very surprising it should remember it and they say well the surprise is measured by the gradient of this loss function right this loss function effectively like let's say you already associate this key with this value um you will have quite zero gradient like okay that's not interesting but um if you don't associate the key with the value yet that's kind of very surprising to you. So you want to learn from it and they say okay let's like you should update your memory and the memory here are just the weights of the neural network by this surprise thing. Um and they add to that and say well if maybe the first token is really surprising but then the next tokens aren't really that surprising anymore but you still want to remember that. So there should be kind of a lag in that. And so we're going to include this pasture prize. And then they're saying like, well, this formulation is similar to gradient descent with momentum. It's like, yeah, that's like you just describe regular gradient descent with momentum by like formulating it as

surprise. Interestingly, they first say, oh, we want to train with surprise and so on. But I I'm not like you can literally describe any gradient descent procedure as this, right? Um to me that is one of the instances where they just formulate stuff more in accordance with this oh it's memory and all of that. Um where uh I feel like they could just say we train the memory to map the keys to the values and we do it we map it by gradient descent with momentum. And look there is this interpretation of gradient descent where you could say it's you know based on surprise. However, they do it the other way around where they say, "No, no, no. Because this is memory, we want really to it to be based on surprise because this is like very human inspired and so on, right? " And um and then it's like, "Oh, but that turns out to be the same as gradient descent and like well yeah so that this is a bit of marketing I feel. " In any case, so the way this works is as I said you are moving across your sequence. Um as you do so you every inference time step right you do a little inner training loop or updating that neural network um that neural memory and um you kind of teach it about the current data point um before you [clears throat] go on and the every step also you pass the current query to that neural network and it will give you back basically an aggregate of the past as it refers to that query. So um that's how you train the memory and that's how you retrieve from the memory right so you have your sequence here and you do attention over the current sequence over what you retrieved from the memory and then over this thing here which they call the persistent memory. So guess what the persistent memory is? It's like, oh, it encapsulates data about the current task. Um, and it's data independent and you can learn it during training. So like this is just parameters. Like this is literally like the fact that it goes here. Um, someone on our Discord said it correctly. This is just prefix tuning effectively, but static, right? So the fact that it goes here has no bearing. It could also be just encapsulated in any of the other parameters of the neural network and it would achieve exactly the same thing. So this to me is again a bit of marketing where they say oh see what you call parameters of a machine learning model. We take just five of them and we call that the persistent task memory that learns about the task which is data independent and you know learned during training. like cool that's you know that we knew that that um and then they say like well look in transformers you also have these things called fully connected layers which are basically similar to attention weights but with data independent parameters. So you yeah like we like this is literally just a fully connected layer. Um this is so they again they start from this a principle of persistent memory and they're like oh well look it's similar to this other thing we already knew where they could have just said um we learn parameters that's it. So this is my first I will rant for one more time here and um I don't want to go through these different variants. They incorporate memory and so on. I feel to me feels that's implementation details and from the experiments it seems like some work here some work there and so on. Um but what I will say is this uh if you have this architecture right and you have this memory and you make a big deal out of the fact that this memory really like the vector valued memory and the matrix valued memory are really not good enough. You need this neural network memory. I disagree. I disagree because this is only valid if you say the way that you store in the memory and retrieve from the memory must be like linear operations or if you specify that your loss function here is specifically um this square distance right here. Um if I can absolutely have a vector valued memory and it will do exactly equivalent to their neural networkbased memory if I

can like the way I store stuff in the memory is arbitrarily nonlinear like this thing here is a super duper neural network and the way I retrieve from the memory is super duper non nonlinear right like it there's no like data is parameters and parameters is data so whether I store things in a simple vector but I make it very complex to retrieve and to store um from the memory or whether um the retrieving and the storing is very simple um but I make the the memory here quite complex those are just equivalent things. And so again, I feel like them making such a big deal and multiple times throughout the paper about all these other models, they only have vector valued memories. We have a neural networkbased memory. To me it feels like that is also kind of like um overdoing it with highlighting the differences here where like in practice if you um if you account for equivalences what you're doing is basically you have a bit more fancy way of retrieving and storing from the memory than the others. Um, and the fact you chose to implement this as the weights of a neural network rather than well ultimately you can simply say that the weights of the neural network and just align them in a vector and and that's what my memory is. So that's it. Um, again I don't necessarily intend to go over the rest of the paper here. Please read it if you're interested. is actually even though I rant about it is a quite a good paper but I just wanted to bring you sort of the core idea here and um I do believe memorizing at test time is an very good idea and we absolutely will need models that do that uh because it's not conceivable to just push context sizes and sort of huristically extend context sizes in the way that we do right now. So something that actively and purposefully learns to memorize at test time and learns to forget at test time is very necessary. So I think this is a really cool direction and the paper is really good and the models uh they perform well and all of that. Uh so that those are my final words here. Thank you and I'll see you around. Bye-bye.