I have a dataset with about 1 000 000 texts where I have computed their sentence embeddings with a language model and stored them in a numpy array.

I wish to compare a new unseen text to all the 1 000 000 pre-computed embeddings and perform cosine similarity to retrieve the most semantic similar document in the corpus.

What is the most efficient to perform this 1-vs-all comparison?

I would thankful for any pointers and feedback!


There are libraries that are specialized in exactly that task, for instance FAISS by Facebook AI Research:

Faiss is a library for efficient similarity search and clustering of dense vectors. It contains algorithms that search in sets of vectors of any size, up to ones that possibly do not fit in RAM. It also contains supporting code for evaluation and parameter tuning. Faiss is written in C++ with complete wrappers for Python/numpy. Some of the most useful algorithms are implemented on the GPU. It is developed by Facebook AI Research.

For cosine similarity, you can use FAISS class IndexFlatIP having normalized the vectors first, as specified in FAISS documentation.

If you want to know more about the techniques behind FAISS, you can have a look at their research paper, and if you want to know more about the fundamentals of similarity search in general, you can have a look at this blog post by Flickr .

Some alternatives to FAISS are Annoy, NMSLib and Yahoo's NGT. You can find a couple of comparisons of these and other libraries here and here.

  • $\begingroup$ Wow, this was exactly what I was looking for. Do you know why you have to normalize the vectors first? To my basic knowledge, cosine similarity considers the angles of the high dimensional vectors, should they not be the same regardless normalization? $\endgroup$ – Isbister Mar 16 '20 at 16:31
  • $\begingroup$ The expression of the cosine distance is the inner product of the normalized vectors. IndexFlatIP does not compute the cosine distance, but the inner product. That's why you need to normalize the vectors. $\endgroup$ – ncasas Mar 16 '20 at 16:36


You could create a pandas dataframe and use a dask function/lambda function to parellize the computation of one vs all at the same time.

If you use dask, you can create partitions and map the response back. In case you use pandas, you can use the apply function and parallelize the computations to a certain extent.


I would suggest also to give a shot to gensim. It's pretty quick compare to other self-written 'top_n retrieval by cosine-similarity' functions.
Save your embeddings in a .txt or .csv file and then load it using the command 'load_word2vecformat'. Once you load the model you can use the function 'similar_by_word' or 'similar_by_vector' to retrive the n closest vectors.

from gensim.models import KeyedVectors

model = KeyedVectors.load_word2vec_format(datapath('model_file'), binary=False)   
top_10 = model.similar_by_word('cat')

Just pay attention when you save the file to include as first line the number of embeddings, 1M in your case, and the vectors dimension. It's just the gensim standard format, the file should look like this:

1000000, 300
'.', -0.0001, -0.0001, ... 
',', -0.0001, -0.0001, ...
'a', -0.0001, -0.0001, ...
  • $\begingroup$ Thanks, I am aware of this approach when using word level embeddings. But I have context embeddings created by the whole sentence! I could store each sentence to a id instead of a word in the standard format embeddings file. But not entirely sure how I could use the similar_by_word("This is a sentence") method, any ideas? $\endgroup$ – Isbister Mar 16 '20 at 8:32
  • $\begingroup$ Sorry, I missed the word 'sentence'. What model are you using to produce them? Bert, elmo or just average of pre-trained embeddings? And are the other pre-computed embedding also coming from the same model? It's relevant cause you might have an alignment problem if the vectors come from different sources, at that point the whole comparison would be meaningless $\endgroup$ – Edoardo Guerriero Mar 16 '20 at 14:44
  • $\begingroup$ No worries. I am using BERT as my language model, all pre-computed embeddings and new embeddings are from the same BERT model. Do you have any additional ideas of best practices here. If so could you update your answer and I perhaps can accept it! Thanks $\endgroup$ – Isbister Mar 16 '20 at 15:57

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.