Going beyond RAG: Extended Mind Transformers - Phoebe Klett

I'm Phoebe. I'm a machine learning engineer at Normal Computing, and I'm really excited to tell you guys about some of our recent research, uh, and in particular Extended Mind Transformers. Allright, so just to briefly cover what we're going to go over in today's talk, uh, we'll introduce the problem, which I think will be quite familiar given the amazing talks which came before mine, and then diveright into the method.

So, uh, what is the retrieval mechanism that Extended Mind Transformers implement? Uh, and then we'll dive into some experiments which give us confidence that these methods are actually performant. After that, we'll get into two of my favorite and I think most compelling features that Extended Mind Transformers enable.

This is a new kind of citation, uh, as well as a new kind of generation paradigm, which is active learning inspired. Uh, and then we'll go over the most important parameters to tune when implementing, uh, EMTs in your applications, and generally how to use them.

Allright. So we pre-train language models, uh, so that they have general knowledge. But as we've been discussing all this conference, that's not enough. We need a lot of application-specific information and a topical, uh, description of the world in order to make these things useful.

Um, I am not going to belabor the two most popular methods, um, which try to load this description into the language model, those being long context and RAG, as I think, yeah, we've heard a lot about those, um, great methods already.

But I'd like to point out that they solve the problem in different ways and thus suffer from different downsides. So long context seeks to extend the context window of the transformer model. So when we train language models, we train them on sequences of a fixed length, and then we're trying to say, well, can we can we extend that so we can include more in the context, more in the prompt during inference time?

Uh, fine-tuning is usually how this is done, and that's awfully expensive. Uh, and more so than that, including all of that context in your prompt can confuse the model with a lot of irrelevant information. Um, and kind of beyond that, just conceptually speaking, it seems a little, like, wasteful,right?

Like, if we're trying to do question answering over a big code base, uh, our queries most usually do not need to reference, like, all of those different function definitions, but just need some subset of them to answer the query correctly.

Um, okay, so this is what RAG tries to do,right? Let's try to subset that information down and just include the most relevant context in our prompt. Um, so what are the issues here? Well, these, these mechanisms which are external to the transformer are kind of, like, necessarily limited by being external to the model.

So we make this choice of what's relevant once and upfront before the generation starts, and we're also making this choice about what's relevant using kind of the least granular representation of that data, and often ones that are disjoint from the way that the model will reason about that data.

Um, kind of also just conceptually, neither of these methods make a difference, uh, or make a distinction between things that should go in memory and things that should be included along with your inference query. And this is more than just aesthetics.

It's actually going to enable us to oh,

it's going to enable us to have these, like, more granular causal citations, uh, and allow the model to retrieve more information when we can tell it's uncertain, kind of actively within the generation.

Allright, so how do we do this? Extended Mind Attention is a very simple edit to the attention mechanism of the transformer. I'm not going to get too much into the math because we don't have a ton of time today, but would love for anyone to check out the paper and let me know what you think.

Um, so but I'll just go over kind of, yeah, from a qualitative perspective how this works. So the model represents data within each decoder layer. Most of the transformers that we're using today are decoder-only transformers. And within each of those decoder layers, the model will represent that data as a key-value pair.

So it actually already has this retrieval mechanism built into the transformer. All we have to do is kind of hack around it. Um, so we pass all of the memory tokens through the model and save off those key-value representations.

And then during generation time, we allow each query token, just like RAG using cosine similarity, to go retrieve a particular number of those memory tokens and attend to them. So this in this picture, these kind of red tokens, red-highlighted tokens are meant to, uh, represent those retrieved tokens.

Uh, again, this actually ends up being a very simple change to the transformer model. What's difficult, uh, is figuring out how to assign position information to those tokens. So this, uh, work is based on research from a couple of years ago, but they needed to fine-tune their model in order to kind of teach the model how to leverage these retrieved tokens.

And that's in large part due to the absolute position embeddings that were popular during that time. So because transformer models are position agnostic, we have to figure out how to kind of tell them, like, okay, this token is position zero, this one is position one, etc., etc.

Um, but due to today's more kind of, like, their softer position embeddings, this allows us to really leverage this method without any further fine-tuning. So in particular, these relative position embeddings that have become popular, and I'll talk about two different methods that we've tested and implemented this on, um, really enable the model to kind of generalize, um, to these retrieved tokens.

The first one, uh, that we tested on is present in all of the LLaMA models. These are the rotary position embeddings. And this generalizes the principle of using kind of, like, an angle between two vectors as a distance metric.

So we kind of take the whole embedding and we rotate kind of two positions at a time. The other one that we implemented, um, this method into is the alibi, uh, linear biases. So these actually aren't position embeddings at all.

It just kind of linearly damps down, uh, information which is further away. And these are, uh, the way that all of the mosaic's MPT models are trained.

Okay, so let's talk about some evaluations. Um, we also just open-sourced a new counterfactual retrieval benchmark, and I'm just going to briefly describe what that benchmark looks like. So this is a long context benchmark. So our input context, these are query answer pairs.

Uh, and the context to answer those questions range from about 2,000 tokens to all the way up to 16,000 tokens. And the again, these are, like, queries. So, like, the question might be, who wrote the song "These Shoes Are Made for Walking?"

and then the corresponding Wikipedia snippet. Um, we wanted to control for facts memorized during pre-training, though, and actually any fine-tuning also. So what we did was we looked up, for instance, in this case, the answer is Lee Hazelwood.

We did a little bit of research. We figured out, okay, well, Terry Allen is a similar songwriter. This is a plausible answer, but it's wrong. We went in and we replaced all the instances of Lee Hazelwood with Terry Allen, and now we asked the model to retrieve this new, you know, not factually correct, but in the sense that we're trying to test whether it's prioritizing what's being provided at inference time.

Um, so now we're asking it to retrieve this Terry Allen answer.

Allright, so how do Extended Mind Transformers stack up? Here, we're comparing it with fine-tuned models as well as the base LLaMA model with, uh, interpolated position embeddings. So we can see here in the green that the base model does a pretty good job extrapolating even, like, many times more.

So this is the model trained up to, like, 2,048 tokens, um, during pre-training. And you can see even up to 8K, it's, like, doing okay. 16K, it really falls off. The position embeddings can't extrapolate that far. The fine-tuned models, you can see, actually perform worse than the Extended Mind model on these shorter inputs.

And this is another data point that suggests that fine-tuning on super long context actually degrades the quality of attention that you get on shorter inputs. And Extended Mind Transformers continue to be competitive with those fine-tuned models all the way up to 16K.

Again, our models are not fine-tuned at all. And in this particular experiment, so what the Extended Mind model sees in context is the query only. So it only sees the, like, who wrote the song "These Shoes Are Made for Walking" and relies heavily on that internal retrieval mechanism to go look up that new information.

In this second experiment, we seed it with a little bit more information in context, uh, using RAG, but again, mostly relying on that, uh, internal mechanism still. Uh, and you can see we're outperforming GPT-4 here now when we combine it with that more information in context as well.

Okay, now we're going to talk about citations. So I think, uh, this would be a topic that lots of you here can empathize with, uh, as AI engineers. I think this is one of the most important things to provide in an application such that people can learn to trust the model outputs.

In fact, you might actually use RAG just to get citations. Um, so with RAG, though, the citations that you get are a little bit kind of, like, post-hoc rationalization. So maybe if, like, the date appears in the output and we knew it was also in the input to the language model, we feel pretty confident that that date is not hallucinated.

Um, but again, this is not really, like, causally related to what information the model used during the generation. Now, with Extended Mind Transformers, we can look up exactly which tokens were retrieved from those memories and used during generation.

So in this example on the top left here, we have the memories. This is a snippet from Wikipedia about one of my favorite mathematicians, Alexander Grothendieck. And the query is, when did he get his French citizenship? And then in the bottom, you can see the completion with the correct date.

I think he got it in 1971. So the blue highlighted tokens here, uh, importantly, the 1971 as well as some of the Alexander Grothendieck tokens, uh, those are the ones that the model retrieved and attended to when generating that 1971 correct token.

And so being able to report that, uh, gives a lot of confidence and also just insight into how the model is using those retrieved tokens.

Okay, we can also use Extended Mind Transformers to reduce hallucinations. So how do we do this? Soright now, we have access to, in the, like, simplest case, just kind of token-level entropy over that output distribution. And if you want to get fancier, we're also doing some Bayesian fine-tuning of language models at normal.

But you can use any uncertainty metric to determine kind of how certain the model is about a generated token. And if we kind of can detect that the model is uncertain about that token, we can regenerate that step using more information from these memories.

Uh, okay, so in the topright here, we can see this is we just set, like, a baseline default number of memories that each query token is allowed to retrieve and attend to. And you can see it wasn't quite enough information, uh, to get this queryright.

So if you remember from the previous slide, the correct answer here is 1971. And you can see we've got 1993 here. So it wasn't enough. We didn't attend to that memory quite enough to get this questionright. And in the bottom example, we allow it to regenerate some subset of those tokens using more information from the cache when we can tell the model was uncertain.

And again, get thisright. So it's kind of, like, kind of a nice intuition for, uh, when the model's uncertain and then, okay, if it's really uncertain, let's go use more information. And also can be more efficient, kind of depending on how the math works out.

Allright, so now I'm going to tell you guys about the most important, uh, parameters to set when using Extended Mind Transformers. So you may have heard of something called STRIDE length before. Uh, and this is, um, a parameter that comes up a lot even just kind of in regular perplexity computations.

So when we compute the memories that we're going to attend to, we pass them through the model and then again, save off these key-value representations that the model saves internally. Um, but again, the models that we're using are trained on this fixed context length.

So we need to kind of pass over them with some STRIDE such that each of those tokens has an appropriate amount of context, um, to generate the representation. So if the STRIDE is smaller, uh, you're going to get more, uh, high-quality representations, but also will require more computations.

Um, so you can kind of tune this. And there are some graphs in the paper as well that kind of represent this trade-off. Um, but this is an important parameter to set when, yeah, generating the memories themselves. Top K is, uh, probably the most important parameter to think about.

So this is the number of key-value pairs or memories that each query token is allowed to retrieve and attend to. Um, when your memory is quite long, kind of the more the better. Um, but again, uh, yeah, this it's kind of should be dynamically set based on how long your memory is.

Um, okay, yeah, so lastly, uh, we want to retrieve as much information as we can from the memory without confusing the model. So making analogy back to kind of putting everything into context, we don't want to just throw everything in there because that will be confusing to the model.

Um, so we have two different regularization techniques that we implement that we have found to be especially effective. Um, the first one is called similarity masking. So again, we we retrieve these tokens, uh, based on similarity with our query token and the key that we are retrieving from.

And so we might say, like, well, if we don't hit some similarity threshold, like, we'll retrieve a lot of them. But then if they, you know, if they're not at least, like, 0.25 similar, then we'll just throw them out.

So we can retrieve and then just mask the ones that end up being less important. Uh, another another important regularization technique, in particular for models that are trained using ROPE, uh, is to eliminate tokens from the memory that correspond to unknown tokens.

So especially if your data is super messy, a lot of the Wikipedia-based benchmarks are, like, really way more messy than I even knew before I started working on this stuff. Uh, they have a lot of, like, just unknown tokens.

And so they're kind of, like, poorly represented by the models often because they're unknown. They end up having a lot of matches with your query tokens, but then they're not actually containing a lot of useful information. Um, so we just eliminate those from the memory before we allow it to start retrieving.

Allright, so we have a whole collection of these models on Hugging Face. All of the, uh, code is on GitHub as well as that data set. Um, and encourage you all to read the paper if you're curious about more of the technical details.

Uh, as I hope you can see here, it's actually pretty easy to use these things. So it's as simple as passing those memories in as inputs, uh, as tokens into the model during instantiation. Um, you can dynamically change them after that as well, but it's the easiest way to do it.

Uh, and then making sure your config is set up correctly.

Allright, so just to conclude here, uh, I hope you all will take away that these new kinds of models, um, impre achieve impressive performance on retrieval tasks. They enable these great new kind of citations. Um, they also enable this new kind of hallucination reduction technique, which is inspired by active learning.

They do not require fine-tuning, unlike kind of long context methods. Uh, and they can be easily run using our open-source models and code.

Thanks so much, and, uh, find me after for questions.