Faiss - Vector Compression with PQ and IVFPQ (in Python)

Key Takeaways

hi welcome to this video we're going to have a look at how we can implement a product quantization index or index pq in physe and we're also going to have a look at how we can implement a composite index using an inverted file list ivf and also pq together to improve our results or our performance even further now we just have a quick look at this if you if you watch the previous video on the logic behind pq you'll recognize this but this is just demonstrating the speed increase that we can get and also the mainly the memory increase which is the important factor in pq and obviously there's a huge increase here so the speed increases is a five to six times increase of just pq and memory decrease is huge it's like a 97 reduction in in memory now i'm not going to go through all of the you know how all that works because we already did in the in the last video um so there is a link to that video in the description so if you'd like to go through that go ahead and go through it but if you just want to see how this works in advice then you know continue and we'll we'll jump straight into it okay so the first thing that we need to do is actually get our data so we're going to be using this uh the sift 1m dataset again in the description there's a link to a notebook where you can download and also load that into your own notebook if you'd like to do that i should say it's a pretty straightforward very popular data set for these sort of examples i use all the time and i'll just show you the sort of shape so xb is um it is like it is the the vectors that we'll be indexing so we're going to index these vectors in our index in face and search through them typically has one million vectors i'm going to stick with 500k because it can just take a little bit of time for everything to load although that being said i don't think it's that bad so let me yeah you know what let's just go with 1 million it's more interesting and then xq i'm just using a single vector here so we just have one vector in there it's c dimensionality for all of our vectors here is 128 which is pretty typical but it's definitely on the lower end for dense vectors so let's let's uh first initialize our pq index so we need to import vice and we'll get our dimensionality which is xb dot shape one now i'm using this slightly different syntax to usual to align better with the pq notation usually you can go with locates dean you can go that's fine it's not really an issue i'm just going to give up case for now to align better to that notation if you are following on from the from the previous video and we're going to set m equals to 8 so this is how many uh sub vectors we have now typically we assert that d is divisible by m but we already know it is so i'm going to skip that but that's something that you you always need to you do need to make sure that d is divided in fact you know i'm going to put it in anyway so let's assert that d is divisible by m okay you just put this in because otherwise we don't we are not splitting our vectors or sub vectors equally and we'll we will not be allowed to do that so we do need to make sure that's always the case and then we're going to set the n bits values so this is the number of bits per per sub quantizer so if you were watching in previous videos we would see this as so the decay star value is equal to 2 to the power of n bits okay so the number of centroids that we have within each sub quantizer and then we we initialize our index so index equals vice index pq d m and n bits okay that's our index ready to go and what we need to do now is actually add our vectors now you just think with with product quantization we are clustering uh our sub quantizers are using clustering so we we do actually need to train it so we can see if an index needs to be trained or not by using this so we are indexed it's trained in this case it's trained as fault so that means we need to train it before we add our vectors so to do that we need to write index train and then we pass it xb run that and this this can take a little bit of time it shouldn't be too long um but with if you start increasing the n bits value which is probably a good idea i mean the end bits of maybe 11 is actually what is recommended in the in the pq paper but i'm not using it here because it takes a long time to to train and build your vectors so i'm avoiding it here but you know if you if you're using this i definitely recommend you know trying those larger and bits of values it's also worth noting that we when we're using composite index of ivf and pq we can't use an n bits value of greater than eight but we'll we'll move on to that anyway uh that has to be a something that they've hard coded into vice for now i think they should be they should be removing that limitation at some point so we have now trained and now we can add our vectors okay so now we want to do is we'll return say the top k most similar vectors will go with k of 100 and then we want to do our search so we're going to go distance i so typically we'd write d here but obviously we already using d so i'm going to write distance equals index.search and then in here we get past our query and we also press k and that will return the the k nearest neighbors in that search so in here if we write shape we see that we have a hundred of those and then the distances here are the actual distances so i the indices of the nearest neighbors distance is the actual distances that they have returned now let's go ahead and and what i'm going to do is i'm going to time it so we can use this time function to essentially run this cell multiple times and get the average time taken to to run it which is quite useful when we're comparing different indexes we'll do index.search and we want xq and k okay now see that that took 3.34 milliseconds so reasonably fast let's um let's compare that to a flat index so we'll do l2 index equals feis index flat l2 and d o d the capital and l2 index.add and we want to add all of our data so it's b and let's time that again so i'm just going to copy this down here and just replace that with l2 index and let's see what we get okay so an average of 19.5 milliseconds so you know already you know this is pretty quick right which is cool now we can check the performance of our index as well now this varies from from all the times i've tested it but we should get something around maybe 50. so 50 recall isn't isn't exactly cutting edge but that is just that's pq that is how is if we want high recoil we can increase n bits to get that uh but it's definitely you know it's a drawback of using uh pq okay so we want so we want to sum for every one value for every value that we find in our index so in our um in the returned indexes for our pq results if they are in the l2 results so we're just saying you know how many of the results returned our pq index were also returned by our l2 index which we are seeing as the 100 recourse l2 index 100 recall and here we're seeing okay how many of those do we match to using our pq index if i in l2 i so i actually need to just copy this pull it down here so when when we use the time it we can't get any variables out from it which is i don't know why but it's a little bit annoying so we will get l2 this and l2 i okay so this time we get 38 recall okay so let's now have a look at the the memory usage of each of these so i'm just going to define a function get to memory so we can more easily compare those and what we'll do is we'll just write our index to file like so index and i'll just call it temp index now because we're going to delete it straight away afterwards we're just writing it to file so we can read see how large the index is and then we delete it again the file size equals os path dot get size to get the size of that index in in bytes i think so we'll do temp dot index and then we don't really we're just using that to decide so we just want to delete that afterwards so like so and then just return the file size okay so let's do that for both of our [Music] both our indexes so we get the memory for the l2 index first uh let me so that's remove not remote okay so we get pretty big that's uh it's half a gigabyte which is massive and then we want to get memory again for our index let's see what we have so now we have we have eight megabytes so that's a massive control that's eight over [Music] five twelve so it's a 90 98.4 reduction in size which is huge uh which really cool now we like i said before we have this composite index where we have both the ivf step and the pq set so we're going to have a look at how we can implement that and we'll see that memory decrease will be a slightly less significant uh because it's a more complex index so they'll be very small still and then the the speed decrease will be absolutely insane so we'll go into that so now we have we have n list which is the number of um the number of cells within our ivf index so if we just have a quick look at what that might look like so imagine these are our pq vectors on a 2d space adding ivf essentially does this so we add all these separating cells called voronoi cells and each of the vectors within each cell is assigned to the vector to the cell centroid now here this you see the the magenta dot at the bottom that's our query vector and it's landed within that cell that you see highlighted now because it's london there we only compare it to those vectors within that cell that are also assigned to the same cell and that's what we're doing with endless we're saying how many of those cells we would like to have now we're going to go with 2048 and it's worth knowing that this must be greater than [Music] greater than or equal to k star which is equal to two to the power of n bits okay so in this case we could go with uh two five six to be fair uh let's can we yeah let's let's go with that for now and see how it goes and then we initialize our index like this so we go index equals vice index iv fpq and in here we need to pass our vectors so in this case we will just when i say vectors the i mean the index that contains our vectors so you typically see this written as quantizer but i'm going to write it vectors because i think it's a little bit less confusing so we'll do vice index flat l2 so we're just going to use these as our starting point but that doesn't mean we're storing these as they are in our in our index they actually get processed using the pq step as we as we already did before so we have effects and d endless m and n bits okay so we now have our index and we can see if it's trained like we did before obviously it will not be because we are still using pq again so we we need to train that and we do the same we did before so train xp after that we add our vectors and we can get our results so index.search xq k okay now we haven't seen how fast that is so let's let's have a look we'll use time it again okay and although the actual search took longer it's because we're looping through you see we have 10 000 loops for this test and we got this number which is insanely small all right so zero point zero point zero zero one that no no that's not right we are on what are we on this 0.1 milliseconds sorry so 0.1 milliseconds what what did we get before it was here so this is using our flat index uh 20 milliseconds if we go up a little further we see 3.3 milliseconds for that search up there so yeah super fast and then if we we check the recall so we'll do we'll use the same as what we used before so where is our sun here like i said this can go up and down um so at 24 this time it's pretty pretty bad but i mean like i said pq isn't isn't the best for recall but what we what we can do is actually because we're using ivf we can increase the number of cells that we search through so you can kind of see that happening here we so using the example from before we start the end pro value of one which is the default increases two three four five and so on and so on and we can increase that all the way up to search all of our empiric values of course that is kind of a waste of time because if you're searching through all of your cells you're basically just doing a flat search with the additional overhead of having a ivf index so it's even slower but we can use that information so we we know that if we set our index.n probe equal to 2048 which is in the number of cells that we the maximum number of cells that we set or did we use two five six uh we know we used two five six so the endless value so we could just button this here if we do that and we do a search again uh xq and k and we go this i do that again and let's come up here get a recall so we see that the maximum recall we're going to get with this instance is 32 now this is pretty low working on it before i was getting sort of 50 52 percent so it's also a case of you know there is some randomness in what you're going to output from your index but of course again we can also increase n bits or and do some other things so we can increase m bits try decreasing m and see how things work but it's also worth noting with ivf pq the end bits value just have to be eight at the moment at least let's just say it's something that's hard coded into files unless you want to go into device code and change that you can remove it if you want now what we want to do is obviously not search through the entire index so we're going to change n probe let me also bring this down i'm going to change n probe to something that is not quite as high so maybe 20. we'll get 33 so that's high enough for us to get our maximum performance but it's gonna be a lot quicker because we're searching through 20 cells or everything 256. let's try 10 32 so increase it a little bit to 12 still 32 16 33 so you know around this area here so 14 13 [Music] so 13 is our sort of optimum empire value where we get the max recall out of that and the quickest time but of course you know which one you which parameters you use completely up to you and depends on your use case as well so that's it for this video thank you very much for watching i hope some useful and i will see you in the next one bye