RAG at scale: production ready GenAI apps with Azure AI Search

Alright, let's get started. Hi everyone, thanks for coming. I'm Pablo, I work in the Azure AI team at Microsoft, and in this session we'll talk about RAG at scale, and in particular I'll focus on the retrieval portion of the RAG pattern.

So planned for this session is we'll do a quick recap, just make sure that we all have the same terms for the same things, and then I'll go through kind of different dimensions of scale, and I'll comment a little bit on how do we tackle this in the context of AI search, what we've learned from doing this, and what are we doing to make it easier to scale these applications.

So kind of by means of a quick recap, when it comes to bringing your own domain knowledge to work together with language models, you effectively you have three options. Either you can do prompt engineering, and you know while it's easy to dismiss, you can go a long way with prompt engineering along, especially these days you know models have longer context and whatnot.

If that's not enough, sometimes the challenge is more along the lines of I want to teach particular patterns, or I want the model to learn jargon of some in some particular vertical domain and things like that, and for that fine-tuning is often a good option.

However, in many, many of these cases, what I want is to have the model work over a set of data that the model didn't see during training. This could be my application data, my company data, company information about my users, or anything like that, and for that the prevailing patternright now is to use retrieval augmented generation.

Effectively, what that means is if you want the model to know facts, then what you do is you kind of separate the reasoning piece of the picture from the knowledge piece of the picture. You lean on the language model for the reasoning capabilities, and you use an external knowledge base to model what you know about a particular domain, the data that you have.

And the way you glue them together is like in principle mechanically very simple. Of course, then it gets complicated because you know life is never that easy, but in principle you have some orchestration component that you know takes the task at hand.

Let's say you have a chat application, and then the user takes a turn and asks the next question. The orchestration component hits some knowledge base seeking for pieces of information that could be used to produce an answer to this question, and then grabs a bunch of candidates, and then you go to the language model and then give it a bunch of instructions, your candidates, and ask the model to produce like an answer to the user question.

That's kind of the essence of the pattern. We all know that in practice it usually takes multiple goes, and there is a lot of tuning in the middle and whatnot, but the fundamentals boil down to that. So just by way of context, like how many of you are doing are creating and working on RAG applications today?

Everyone, okay, excellent. So with that backdrop, what I wanted to do is talk about what are the pressure points when you scale these applications. One thing that has been fascinating to see is you know we had the opportunity to be involved in this space from very early, like Azure OpenAI has been there from you know the early days of scaled language models, and one thing we saw was all of I'd say last year, 2023, everybody built a prototype of something to kind of learn and figure out what could be done with this technology.

And the interesting shift to this year, to 2024, it's been these applications are going to production, and when you go to production you go from oh this demo is really cool to all these users are you know are using it at the same time, and they want more, they want more data, they want more faster answers and whatnot.

So the result of that is that you know when before we could focus only on figuring out the interaction model and the applicability of this technology, now you know elements of scale also play a role. And scale kind of takes multiple kind of flavors.

Like you these things tend to scale in volume because one of the things that happens is when your application works well, users come back or the leadership of the organization comes back and says, "Let's put all the data there," and then you have to deal with it.

Also the rate of change of the data kind of increases, and the query load increases as well because you know more people are using the stuff. Also workflows tend to get more complicated. At first you go like oh like how complicated it can be, I'll take the question, search the thing, and then send it to the model.

Well, it turns out that sometimes it works, but often it doesn't. So you end up doing these multi-step workflows that hit the retrieval system and the language model multiple times, and that taxes all the systems, and they all have to scale.

Also in the kind of in the spirit of now let's put all the data there, now you have to deal with more data types, different kinds, different data sources, and whatnot. So let's cover each of these dimensions of scale in detail, and what I'll do is I'll cover it in the context of Azure AI Search, because that's where I work mostly, and because it's you know the way we think about this in Azure is we want to produce a system that encompasses the entire retrieval problem.

Of course, it has this kind of vector database capabilities, the vector-based retrieval emerged as a very useful solution in many contexts, so we have full support for that, but we also brought into it like years and years of Microsoft experience in retrieval systems in the more general sense, and we integrate them together so you don't have to connect a bunch of parts to have a proper high-quality retrieval system.

It all comes integrated, and of course you know we integrate it into the rest of Azure so it's easy to pull data in, to connect to the other data sources, and whatnot. And you know Azure is a place that is used to build some of the largest applications out there, so all the kind of enterprise readiness comes pre-builting, you know from security to compliance to all these things that you don't want to deal with directly.

You want to kind of build on a platform that is already taken care of. So while there are multiple moving parts to retrieval systems, what we've seen the last 18 months or so since the emergence of kind of RAG patterns at scale is vector retrieval is an important part of the solution, and you know you can see why.

Like the interesting part about vector retrieval is like it's incredibly effective at getting this kind of soft conceptual similarity and put it to workright away. So in Azure AI Search, we built a system that has kind of complete feature set when it comes to vector search, including fast approximate nearest neighbor search, and as well as also exhaustive search.

Sometimes you want precise search either to create baselines, to know how kind of your recall performance is looking, and things like that. Also often applications need to combine vector search with the rest of the queries. You need to filter things, you need to slice and dice, like you know filter, select the columns you want to project, and things like that, like effectively treat it like a database that also does retrieval.

And we also see multiple scenarios where the documents have multiple queries, maybe for different parts of the content or for different data types that use different embeddings, and sometimes the queries need multiple vectors. So we try to make it as this all these specific needs come up as you're building an application, you will find answers to all of them directly in Azure Search.

Let me show you this in action.

So let me start with a very simple example. So what I'll do is I'm just here I have a small Jupyter notebook, and what I'll do is I'll connect to I'll kind of set up a default credential point at my Azure Search service, and I'll create an index from scratch.

So this is all like all it takes to create an index once you have a service provisioned. So in this case, I'll create a few fields. I'll create a categorical field like this serves as metadata that you want sometimes to attach to be attached to this.

I'll create a text field. Sometimes you want to mix text and vectors, and I'll talk a little bit more about this later, and I'll create a vector field. In this case, this is a toy example, so dimension is three.

Of course, that's not very useful. You know in practice this is kind of in the hundreds or maybe thousands of dimensions. I'll also say what strategy you want the system to use for vector search. In this case, I'm using HNSW, which is a well-known graph-based algorithm for indexing vectors.

So when I run this, and I hope like the network was a little wonky, so oh there perfect. So now I have an index, I'm going to get a kind of a client to it, and then I'm going to index data, and you can see here this is a very simple case.

These are my vectors, this is my full text, and these are my categorical data bits. So you can let us do all the ingestion for you, or you can push data into the index. In this case, I'm pushing data explicitly into the index, and once you have data, if you look at here, like we have some vectors and we have some categories.

So you can first you can just search, and because it's in this case I'm indexing vectors, when you search, you search with a vector. So I can search for that and say these are the two closest to the reference vector that I gave, and I can do a few things kind of incrementally.

Like for example, if I want to combine text and keyword search, I can also add search textright here. I'm going to go hello, and now I'm searching for both the text and the vectors, and then we fuse the results and rank appropriately.

And often in applications, you also need to filter stuff. You can see here I have A's and B's in categories, so I can like write filters. For example, here I'll say category equals A, and then you know I only get the A's.

This is a trivial example, but of course you can do full filter expressions and ANDs and ORs and ranges and whatnot, and the filters are fast. So even if you have hundreds of millions of documents, these are not a problem.

So in that one if you had k nearest neighbors equals two, how come they had three results?

Oh because a great question. Because so what I was telling this the system here is from the vectors retrieve two candidates, but then I was like I also told it go to the keyword side and retrieve a bunch of candidates from there and then fuse them.

So only two of them were from vectors, but the fusion of the two selected three. I actually like if I can make this a larger number, and that sometimes is useful to get more candidates, and then still I can say as a result of that get the keyword ones too, and then rank the top end.

So basically separate how many candidates you want from how many you want to return.

So those are the basics of a vector engine, but there are a few key elements to consider when you are actually building an application in production. Probably the most salient one is quality. Like in the end, your application works well for your users when they ask a question and they get the answer they're looking for, and that is highly influenced by whether your retrieval system actually found the bits of information that you wanted to produce that answer.

In AI search, the way we do this is we do what kind of most kind of more sophisticated search engines do, where you use a two-stage retrieval system. First stage is recall-oriented and uses vectors and keywords and kind of all these recall-oriented tricks to produce as many candidates as we can find, and then the second stage re-ranks those candidates using like a larger model that you maybe let's say you have 100 million vectors in your database.

You want to use something fast to go from 100 million to a small set, but then once you have a small set, you can afford to run a larger more sophisticated language sorry ranking model to create better quality ranking.

So that's what we're doing in the L2 stage. So when you turn this on, you effectively get better results, and what you can see here there is a link at the bottom. Happy to share it later with more details on the numbers, but you can see that so from the left this is what you get when you just do keyword search using BM25, which is a well-known scoring approach.

Second one is only vectors using AIDA, OpenAI AIDA vectors. Third one is using fusion, combining vectors and keywords, and the fourth one is using fusion plus a re-ranking step, where across the board we see better results just out of the box when re-ranking is enabled.

Yes?

So the re-ranking is the semantic re-ranking. Is that another cosine similarity style thing, or how does that actually work?

Great question. The question is, is it a cosine similarity thing? No. So the thing by encoders are useful when because you can do you can encode all your documents as vectors and then the query as a vector, and then you can evaluate them fast because you're only comparing similarity, but that means that at no point the model sees the document and the query at the same time.

So the re-ranker is this type of rankers are often called cross-encoders, and these are transformer models that you feed them the document and the query, and you ask them to predict a label of the correspondence of the query to the document.

So they are much better positioned to produce a high-quality result, but they're at inference time, or rather at query time, you're running an inference step. So you can do it only on a smaller set. You couldn't do it on the entire data set.

It wouldn't be practical, at least not for interactive performance applications.

And speed?

Speed? This is like 100 milliseconds give or take for a model like this. Depends on the cross-encoder that you use. Ours our trade-off between you know make it fast enough, but make it very high quality. We landed on 100 milliseconds ballpark, and we found that that works well in terms of interactive performance because the majority of the latency ends up hidden on the LLM call, and you still get very high-quality results.

The other thing I won't drill into this, but often the other dimension of getting quality out of the system is to narrow the data set. If you know discrete metadata elements that can help you narrow the data set, that's the most effective way.

Then do all the ranking tricks on top of the resulting set, but first scoping is a very effective way to get quality up. Yes question.

So my question is, do you extract the keywords from the top, or like how do you set the keywords?

For the keyword part of keyword search?

Yeah.

Oh so if you look at let me show you this example. So if you look at what I did here is you give us keywords, you give us text, and then you give us a vector as well, or you just give us the text and we'll turn it into a vector too.

So we need the original text. The vector alone is not enough because we can't go back to a keyword.

No so I mean if you have a RAG application, let's say you have a chatbot.

Yeah.

How can you know ahead of time what kind of keywords will be added onto that specific prompt that the user is given?

Oh so usually in a RAG app, the question is what of the conversation so far do you want to use for search? Is that the question?

No I understand, but my question is so if you want to use because in your metrics you showed that the fused search has higher performance. So how do you like for vector search it's easy,right? You vectorize the prompt and then search.

Yeah.

But if you want to add like search and vector field.

Yeah.

Is there a way you can extract like probably use it rather than to extract the potential keywords for your search because you cannot pre-fix those keywords for every potential chat or prompt that the user can have.

Yeah I maybe we should talk after the talk because I'm not sure I understand the question. Like usually like you in the index you have the text and the vectors, and then from the user you extract a candidate search for that.

This is before you send it to the LLM, but maybe we should chatright after the talk on the details. So let me go back here. So the other dimension of scaling is just how many vectors and how much content you can fit in one of your indexes, and for that in AI search like one thing we learned last year at the beginning of this year is that it went from everybody had tiny indexes to everybody put all their data in these systems.

So we accommodated for that by significantly increasing the limits. Like for most of the search services, you will get anywhere between 10 and 12 times the vector density on the same SKUs, and we didn't change prices or anything, just so you can build larger applications on the same setup.

With these new limits, you can build multi-billion vector apps by just provisioning a service and uploading the data, and it's surprisingly straightforward. Like you know a year ago, the billion data set vectors was the thing there was a curiosity used for benchmarking, and now you can just create an index and upload them, which is very impressive to see.

I'm not going to drain the slide. I put it here for reference. These are kind of all the new limits that we have. They're significantly higher than the ones we had before.

One of the things that has been exciting for us to watch grow is you know among our customers, one of them is OpenAI themselves, where they have a lot of RAG workloads. Like when you have files in ChatGPT or when you use the assistance API, all of those you can create these vector stores inside their system, and all of those are backed by AI search.

And when we increase these limits, one of the things they did is they increased the limits they give to their users by 500 times. So it's been impressive to see how fast they grow, and it's been fun to kind of see kind of first from up close how a system can scale like this big and still be fast and responsive and whatnot.

So finally, the other thing you can do is sometimes the limits higher limits are enough. Sometimes you want to push even more data into it. So the other thing we've been working on is quantization, where you can use narrower types.

Like instead of using full floats, you can use ints, like int8, and they just simply use less space at the trade-off for a little bit of quality. And interestingly, you can even do single-bit quantization, and I confess that when people said, "Hey, like we're going to do metrics for single-bit," I felt people were wasting their time, but it actually works.

It works surprisingly well. Our evaluations show that they are still for some models in the mid low to mid 90% of the original performance, and other companies have seen the same thing. Like for example, this is an evaluation from Cohere, but separate company, and they also see like about almost 95% of the original precision is preserved when using bit encoding, and you can get the precision remaining precision back by re-ranking when you're done.

So surprising that it works, but you go for float 32 to one bit, that's 32x the vector density, and it's faster because you just do humming distance, much faster than computing on a smaller number of bits instead of cosine similarity or something like that on a wider set.

So because of this, we now support all these types in AI search as well.

I'm going to skip this slide in the interest of time. So and you can do quantization yourself, or you can just enable it, and we will do quantization for you, and if we do quantization for you, we will also store the original precision data there, which means we can do oversampling where we query at the quantized kind of compressed version of the vectors, but we have the full precision stashed on the side.

So later we can re-rank at full precision, and so you can effectively choose between you want a highly compressed index that is a little lower quality but larger, or it's a little slower but larger, or you just don't compress it and then you get the quality up.

So effectively, you can choose any of the three, and it's effectively up to you what you want to prioritize, but we give you control for all three dimensions.

And then the last thing I wanted to touch on is that the other challenge is you have to keep adding data sources that you bring into these RAG systems, and each of them you connect to them differently. You enumerate changes differently, and that's just not where you want to spend your time.

So we have this ingestion system that includes integrated vectorization as part of AI search, where if the data is in Azure, whether it's in blob storage or OneLake or Cosmos DB, we will connect, deal with all the security and all of that.

We will automatically track changes. So it's not a one-shot thing, but as the data changes, we'll pick up only the changes and we'll only process the changes. So the cost is also incremental. You don't pay as you update the stuff for the entire set, and then we'll deal with all file formats, you know PDFs, Office documents, images, unpack the nested formats and whatnot, do chunking, do vectorization, and land it on an index all in one go.

And this is a like industrial strength pipeline that you set up once, and then it continuously runs after that. As your data changes, it reflects the changes. So you can focus on the RAG stack, you know the workflow, how you query the system, but you don't have to think about how the data makes it there.

If the data is anywhere in Azure, we'll index it and create like an index that follows the original data automatically. Allright, and with that, I know I raced through this content in these 20 minutes. I'll be hanging out outside if anybody wants to chat or have questions, and I would encourage you to go try AI search today.

Here's a link to the starting point, and Azure subscriptions include a free instance of AI search, so you can even give it a shot in a minute without having to pay for any of this stuff. With that, thank you.