Going beyond RAG: Extended Mind Transformers

[upbeat music] 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.

All right, so just to briefly cover what we're gonna 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, uh, and then dive right 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.

All right. 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 they suffer from different downsides. So long context seeks to extend the context window of the transformer model.

So 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 query is most usually, does not need to reference, like, all of those different function definitions, but just needs 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 up front 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 gonna enable us to... Oh.

It's gonna 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.

All right, so how do we do this? Extended Mind Attention is a very simple edit to the attention mechanism of the transformer. I'm not gonna 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, was figuring out how to assign position information to those tokens. So this, uh, work is based on research from a couple 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, uh, 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 to position one," et cetera, et cetera.

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.

These actually aren't positioning 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 gonna briefly describe what that benchmark looks like. So this is a long context benchmark, so our input context is our query answer pairs.

Uh, and the context to answer those questions range from about two thousand tokens to all the way up to sixteen thousand 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 Hazlewood.

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 Hazlewood with Terry Allen, and now we ask 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.

All right. 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, ex- 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 was a model trained up to, like, twenty forty-eight tokens, uh, during pre-training.

And you can see even up to eight K, it's, like, doing okay. Sixteen K, it really falls off. The position embeddings can't extrapolate that far. The fine-tune 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 sixteen K.

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 gonna talk about citations. So I think, uh, this will 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, a 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 tho- 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 nineteen seventy-one. So the blue highlighted tokens here, uh, importantly the nineteen seventy-one, as well as some of the Alexander Grothendieck tokens, uh, those are the ones that the model retrieved and attended to when generating that nineteen seventy-one 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? So right now we have access to, in the, like, simplest case, just kind of token level entropy over that output distribution.

And if you wanted 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, so okay, so in the top right here, we can see, uh, 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 query right.

So if you remember from the previous slide, the correct answer here is nineteen seventy-one. And you can see we've got nineteen ninety-three here. So wasn't enough. [laughs] We didn't attend to that memory quite enough to get this question right.

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 is uncertain.

And again, got this right. 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.

All right, so now I'm gonna 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 gonna get more, uh, high-quality representations, but also it 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 d- 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 is kind of, should be dynamically set based on how long your memory is.

Um, okay, yeah. So lastly, uh, we wanna 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 wanna just throw everything in there 'cause 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 .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, [laughs]

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 n- not actually containing a lot of useful information.

Um, so we just eliminate those from the memory before we allow it to start retrieving.

All right, so we have a whole collection of these models on Hugging Face. All of the, uh, code is on GitHub as well as that dataset. 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.

All right, so just to conclude here, uh, I hope you all will take away that these new kinds of models, um, impre- in- 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. [audience applauding] [upbeat music]