I have an embeddings matrix of a large no:of items - of around 100k, with each embedding vector length of 100. So a matrix of size 100k x 100;

From this, I am trying to get the nearest neighbors for each item using cosine similarity. I have tried following approaches to do that:

  1. Using the cosine_similarity function from sklearn on the whole matrix and finding the index of top k values in each array. But I am running out of memory when calculating topK in each array

  2. Using Pandas Dataframe apply function, on one item at a time and then getting top k from that

    similarity = df[embField].apply(lambda x: cosine_similarity(v1, x))
    nearestItemsIndex = similarity.sort_values(ascending=False).head(topK)
    nearestItems = df[itemField].ix[nearestItemsIndex.index]

But this approach is taking around 6-7 secs per item, and is not really scalable.

As this should be a common case in recommendation systems, I am guessing there should be some existing algo to solve this on large data. But unfortunately I couldn't find it. Would be great if someone can help me point to any such algo.

  • $\begingroup$ Normalize the vectors to be unit vectors. Use a k-d tree for nearest neighbor search. $\endgroup$
    – user18764
    Mar 27, 2018 at 14:36
  • $\begingroup$ This thread may be of interest: How to overcome the computational cost of the KNN algorithm? $\endgroup$
    – user20160
    Sep 27, 2019 at 10:23
  • $\begingroup$ I don’t think you can. Once you use cosine similarity you lose the magnitude. So two points can have have 0 angel, meaning cosine similarity of 1, but can be very far away $\endgroup$
    – Nathan B
    Mar 29 at 13:32
  • $\begingroup$ @NathanB cosine similarity is a common measure of nearness in ML. Often, the vectors are unit-normalized while training a model, so magnitude doesn't play a factor. $\endgroup$ Mar 29 at 15:27
  • $\begingroup$ Only if the embedding vectors are unit-normalized which means that they are on some hyper-circle, only then the nearest-neighbours is equal to the "closest" cosine-similiary embedding-vectors. $\endgroup$
    – Nathan B
    Mar 31 at 9:46

3 Answers 3


Actually, we can use cosine similarity in knn via sklearn.

The source code is here.

This works for me:

model = NearestNeighbors(n_neighbors=n_neighbor,

one weak point is sorting, and second is creating new collection

so you could just make another collection (list) and then keep topK items until some point this way instead of sorting everytnig and allocate new portion of memory you could do it more robust use this https://stackoverflow.com/questions/27672494/how-to-use-bisect-insort-left-with-a-key

# import my_insort_left

topK_items = [] # (index, value)
for i, row in df.iteritems():
    res = cos_sim(v1, row)
    my_insort_left(topK_items, (i, res), keyfunc=lambda v: v[1])
# topK_items is the result
nearest_items = [k for k, v in topK_items]

*keep in mind that I wrote it without testing so it may have some tiny bugs


A Nearest Neighbours model is fairly fast to build, because the algorithm uses the triangle inequality. Sadly, Scikit-Learn's ball tree does not support cosine distances, so you will end up with a KDTree, which is less efficient for high-dimensional data.

from sklearn.neighbors import NearestNeighbors

embeddings = get_embeddings(words)
tree = NearestNeighbors(
    n_neighbors=30, algorithm='ball_tree',
  • 1
    $\begingroup$ I dont see SciKitLearn's KDTree or BallTree algorithms supporting cosine metric. Have you got this working with KDTree? $\endgroup$
    – Vineet
    Dec 10, 2019 at 16:24
  • 2
    $\begingroup$ sorted(sklearn.neighbors.VALID_METRICS['ball_tree']) will print the valid metrics applicable to NearestNeighbors, 'cosine' is not among them $\endgroup$ Aug 28, 2020 at 16:17

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.