# The Memory Problem, Baseten | Compile 26

Channel: Cursor
Video: https://www.youtube.com/watch?v=I8YnwUV2C9w
Duration: 13 min
Language: English (United States)
Words: 2067
Transcript page: https://viewrankai.com/tools/youtube-transcript/I8YnwUV2C9w

---

[0:00] We're part of the research team at Baseten, and one of the big problems we focused on a lot is the memory problem. And a lot of the stuff that Ben talked about was really useful motivation for the stuff we're going to talk about. And hopefully, we are one step closer to this kind of lossless retrieval of the relevant context at inference time, which is a lot of what our research has

[0:20] been recently. So, in terms of memory, I guess we see there broadly being two main types. You've obviously got your relatively lossless memory, which is your full KVK, and as we know, that kind of scales linearly in terms of the memory requirement, but this does become a very clear bottleneck when you are trying to do long-horizon agentic use

[0:45] cases. And we've then got highly compressed versions of memory, and this includes things like fine-tuning, RAG systems, and writing to markdown style files in the form of natural language compaction. What is missing is this intermediate memory layer, which we really believe is important to get these long-horizon agentic use cases working. And taking a step back, I guess, what are the key components to getting a

[1:17] reliable long-horizon agent to do useful things? We see there being three main components. So, we have to have an effective version of compressed memory.

[1:31] And for us, that looks like compression in the KVK space. The second piece is being able to use memory in the loop. And the main idea here is if memory can be used in a loop repetitively, this is how you're effectively going to get these really long-horizon tasks to be able to use the memory effectively. And then the third and final component is some sort of mechanism to project the memory back into the weights. And the idea here is, how do we turn these experiences

[2:07] into some sort of durable knowledge? As mentioned, we really believe compression of the KVK space is the thing that is emerging as the way to be able to do this effectively, and we see these kind of existing in a two-by-two matrix in terms of the methods that exist. So,

[2:25] on one axis we've got selection versus synthesis. And on the other axis, we've got per context or amortizing this cost in a way. So per context selection involves looking at the full KVK space and having some sort of scoring and being able to choose which parts or which parts to evict of the KVK so that

[2:58] that can be used for future calls. In terms of amortization, what we mean by here is paying the cost of training upfront, and that cost is therefore amortized on all the inference time contexts that are going to be used for that compressed KVK space. In terms of per-context synthesis, I guess the key distinction here that we're saying between selection and synthesis, selection is still represents some

[3:24] subset of the original KVK space. But when we say synthesis, we are saying some learned representation, often through backprop into the KVK. And a really good example of this is the Cartridges paper, where we see an effective representation in KVK space that is done at context time. Our method of Still exists in this quadrant here, and Charlie will

[3:47] motivate that a little bit later. And just for the sake of completeness, I guess at pre-training time, we see methods that exist both in selection and synthesis. So DeepSeek, as Ben has already talked about, DeepSeek's sparse attention is a great example of selection of appropriate parts of the KVK space. And then synthesis, again, with their compressed sparse attention and heavily compressed attention. There are components in there that is synthesizing parts of the KVK. So, I'll pass over to Charlie now to talk about our method, which

[4:19] exists in this quadrant.

[4:22] So, I guess a big motivation for this work was this idea of amortization. And so a lot of people are familiar with the idea of sparse autoencoders and sparse coding in general. So if I have some sort of input that might be neural activity in the brain or a dense vector in an LLM, very dense, there's a lot of superposition,

[4:44] a lot of stuff going on. The traditional sparse coding would run this iterative algorithm every time I fed in a new input to try and figure out what the latents were. So this might eventually figure out that the active latents are dog and ocean, if it's a picture of a dog by the beach. Then sparse autoencoders came along, and what they said was, okay, instead of running this iterative optimization for every new input that I feed in, what

[5:09] I'm going to do is learn a general function. So an F I can apply to my X, which should, across all my training data, WX + B, learn to output roughly the same correct codes. And so the idea here is you're amortizing all of your inference compute over training, and essentially what you're doing

[5:33] is learning the optimizer.

[5:38] And so basically, we just said we're going to apply this same idea to our problem, which is compressing the KVK. So as Moody said, we have this full KVK most of the time with cartridges and selection methods and all this sort of stuff. We're running an expensive inference time algorithm to optimize what the

[5:56] compressed KVK should be. And we said, okay, instead of doing that, what would this black box function equivalent look like? So we can take the full KV cache and output a compressed KV cache. And so that was the motivation. When you get down to the weeds of what this should look like, essentially, we have a variable-sized KV cache

[6:22] because it can be a variable number of tokens, and we want to compress it. So we already know we need to use some form of attention or transformer. Let's say I have

[6:34] two queries, so lowercase t equals two. I'm going to learn queries which are able to ask general questions about the KV cache that I have. So this KV cache might be any number of tokens long. So let's say, in this case, I have five, so capital T equals five. And what I'm going to do is I'm going to take my queries that are going to be learned by a still model, and I'm going to cross-attend

[7:02] and learn to ask questions of every single token in the KV cache. So here you can imagine this is a big X, and we've kind of just concatenated our keys and values. Cool. And so we get a very familiar formula. So I'm going to have some attention score, which is just softmax of my Qs applied to my KXs, and KX here is simply what I get when I take my concatenated Ks and Vs and apply my

[7:34] learned key matrix. Same with V.

[7:40] Cool. So I get this, divide by a spirit of dimension, and so I have an attention score over my queries. Great. Basically, then we can get the output by simply taking these attention scores, matrix multiplication with VX, and then having some sort of output projection. So this is the attention that everyone knows and loves. This to actually be the identity. And so at the beginning of training, all we're doing is taking

[8:08] mixtures of existing Ks and Vs and outputting that. And then essentially, the loss function here is the KL of the model with the full KV and the model with the compressed KV. And we're trying to get the model to output the same stuff with the compressed KV, and that's a really nice, kind of very natural

[8:31] next token loss function that we use. But I'll pass over to Harry to talk about some specific implementation stuff. Cool. Charlie basically here was talking about doing a single compaction, and what we want to use this for is actually infinite context. And so I'm going to talk about some problems that we encountered and how we solved them for doing iterative

[8:54] compaction. So just write this off.

[9:01] So the naive implementation, if we want to use this kind of compactor iteratively, is so we'll have a KV cache block, and we compress that into a smaller block. And then let's just say, okay, now we're going to start generating from this compressed KV cache, we're going to generate a new block, and then we're going to compress that,

[9:26] and then we're going to keep going. So this is the naive implementation, and it works to some extent. But basically, we find that if we train with just single-pass compaction and then try to use this iteratively, it doesn't work for many iterations. After a few iterations, it starts to collapse. And what we found is that a key thing to stabilize this, was instead of when we're doing the KL divergence, the typical training objective that we use is to have the full context and the compacted context, and then have the next KV cache, and compute KL divergence

[10:02] on this next cache between this full context and compacted cache. So what we found is that instead of computing on this next block, if we instead have the subsequent block here, and then we compute KL divergence of the model with this full context and this block here, so we compare this full section to this section here, and we compute

[10:28] KL divergence on this subsequent block. So that was one technique which we used to essentially... And now we can do 16 or even 32 iterations of compaction, and it's still retaining quite high accuracy. So the next thing I wanted to talk about was what is actually happening to this compacted cache as we do this iterative process, and what really is this compacted cache? So something which I like to talk about is that,

[10:55] essentially, a KV cache, I'm going to draw this as... Just first of all, draw the keys, and then I'm going to draw the values as being transposed. So if we think about what happens with a single query, when we have a single query coming in, here is our Q vector, or maybe I'll make it a lowercase q. So your query vector comes in, you matrix multiply with your set of keys, and you get an initial attention scores. You apply a softmax to these attention

[11:26] scores, so we apply a non-linearity. And then we have these attention post-softmax scores. Now we matrix multiply with the values, and we get the final outputs here, and we might also have a output projection. Now, the key thing is, you'll notice that this is actually basically exactly the same procedure as a normal MLP. This is the up projection of an MLP, and this is a down projection of an MLP,

[11:52] and we have a non-linearity in between. So really what's going on here is a compacted KV cache is an MLP. We have constructed from the context an MLP. And this is basically a new form of producing weights. The traditional way in which we produce refined weights is through gradient descent. But basically, what we're doing here is producing weights through compaction. And the potential advantage of this is that in-context learning is incredibly sample efficient. You just say one thing in context, and the model obviously remembers that fact because it can see it in this

[12:27] context. Whereas traditional gradient descent is very sample inefficient. It takes many examples in order to learn something. But if we now have a mechanism to produce weights directly from the KV cache, this allows us to just have said something once in the KV cache and then compress it. And then if we can combine that with the weights or just retain it as it already is a set of weights, essentially, then we basically have much more sample efficient learning. That's where we're interested in going in the future and basically trying to do

[13:00] gradient descent and compaction simultaneously. So we might want to do compaction to produce things in the first place and then also do gradient descent into that same block, and that is something which we're going to experiment with. Cool.

---

About this transcript
Read from YouTube's own caption track and laid out by ViewRank AI (https://viewrankai.com).
ViewRank AI finds the videos already beating a creator's own average on Instagram, TikTok and YouTube Shorts, transcribes them from the audio itself in more than 60 languages, and turns what worked into new ideas and scripts.
Free transcript tools, no account needed: https://viewrankai.com/tools
How to read any video this way: https://viewrankai.com/llms.txt