Package: fast-ivf
FastIVFEfficient implementation of IVF Index with numpy and numbaInstallationInstall numpy and numba from conda to use intel mkl libraries for linear algebra operationsTo install package from the source code runpip install .To install from pip runpip install fast-ivfYou may need to install tensorflow>=2.13, seeCompressedFastIVFfor detailscode tested with python==3.11see notebooktest-indexfor Index usage examplessee notebooktest-kmeansfor K-means usage examplesFeatures / limitationsThis is an experimental code which heavily relies on numba and numpy and may contain bugsIVF centroids are estimated with custom mini batch kmeans implementationMiniBatchKMeansis used to estimate centroids of standard Inverted IndexSubspaceMiniBatchKMeansis used to estimate centroids of Product Quantization IndexK-means implementations support only l2 or cosine distancesAll indices currently support only cosine distanceResults on custom benchmark dataResources restricted toOMP_NUM_THREADS=MKL_NUM_THREADS=OPENBLAS_NUM_THREADS=12which was consuming 100% in our case for fast-ivf and faissTrain vectors: internal ~900k vectors of dim=1024, normalized to unit lengthTest vectors: same but 40k vectorsHyperparams: nprobe=10, ratio_threshold=0.5, no re-scoring is used for approximated indices (for mini-batch kmeans we use repository defaults),
for CompressionFastIVF we use compression_ndim=128 (which gives 8 times compression ratio)We measure recall@10, as function which checks ifexact_i is in top_indices[:10]for each test query, then we
average the results over all test vectorsFor faiss I used similar parameters for nlist, m, nbits etcReported time is computed from average of 5 runs, divided by 40k to get the time per single queryAs we use numba internally, each Fast-Index is initialized with warmup call to compile the codeNote: CompressedFastIVF requires to train small neural network to compress embeddings to lower dimensionality, which increases the index build timeFor both libraries each search() call was consuming all 40k vectors, to fully utilize all vectorizationIndexRecall@10Query Time (ms)ParamsFastIVF0.9640.100nlist=1024, nprobe=10, ratio_threshold=0.5Faiss IVF0.9681.000nlist=1024, nprobe=10FastIVFPQ0.8020.100nlist=1024, nprobe=10, ratio_threshold=0.5, pq_num_subvectors=32, pq_num_centroids=128Faiss IVFPQ0.8640.220nlist=1024, nprobe=10, m=32, nbits=7CompressedFastIVF0.9330.050nlist=1024, nprobe=10, ratio_threshold=0.5, compression_ndim=128CompressedFastIVF0.8890.040nlist=1024, nprobe=10, ratio_threshold=0.5, compression_ndim=64Custom mini batch k-means implementationEfficient mini-batch kmeans implementations with numba and numpyfromfast_ivf.kmeansimportMiniBatchKMeansimportnumpyasnpkmeans=MiniBatchKMeans(num_centroids=16,batch_size=32,metric="l2")data=np.random.rand(5000,64)kmeans.train(data)kmeans.add(data)labels=kmeans.predict(data)Efficient mini-batch kmeans implementations to train product quantization centroidsfromfast_ivf.kmeansimportSubvectorsMiniBatchKMeansimportnumpyasnpkmeans=SubvectorsMiniBatchKMeans(num_centroids=16,num_subvectors=8,batch_size=32,metric="l2")data=np.random.rand(5000,64)kmeans.train(data)kmeans.add(data)labels=kmeans.predict(data)FastIVFSimilar tofaiss.IndexIVFFlat( faiss.IndexFlatIP(d), d, nlist, faiss.METRIC_INNER_PRODUCT)fromfast_ivfimportFastIVFfromfast_ivf.coreimportnormalizeimportnumpyasnpnlist=1024train_embeddings=normalize(np.random.rand(10000,512).astype(np.float32))index=FastIVF(512,nlist=nlist)index.train(train_embeddings)index.nprobe=10# greedy skip voronoi cells which are having score smaller than 0.5 of the largest score# higher values lead to faster search but less accurateindex.ratio_threshold=0.5test_embeddings=normalize(np.random.rand(100,512).astype(np.float32))distances,indices=index.search(test_embeddings,k=100)FastIVFPQSimilar tofaiss_index = faiss.IndexIVFPQ(faiss.IndexFlatIP(d), d, nlist, m, n_bits)fromfast_ivfimportFastIVFPQnlist=1024# pq_num_centroids = 2 ** n_bits# pq_num_subvectors = mindex=FastIVFPQ(512,nlist=nlist,pq_num_centroids=64,pq_num_subvectors=32)index.train(train_embeddings)index.nprobe=10index.ratio_threshold=0.5distances,indices=index.search(test_embeddings,k=100)# compute exact scores for top 100 results, this is slower but more accuratedistances,indices=index.search(test_embeddings,k=100,rescore=True)# calibrate scores by fitting a linear regression model to N=20 exact scores, if -1 then all scores are exactly computedindex.rescore_num_samples=20distances,indices=index.search(test_embeddings,k=100,rescore=True)CompressedFastIVFTrains keras autoencoder to compress embeddings to lower dimensionalityfromfast_ivfimportCompressedFastIVFnlist=1024index=CompressedFastIVF(512,nlist=nlist,compression_ndim=128)index.train(train_embeddings)index.nprobe=10index.ratio_threshold=0.5distances,indices=index.search(test_embeddings,k=100)# compute exact scores for top 100 results, this is slower but more accuratedistances,indices=index.search(test_embeddings,k=100,rescore=True)