search

rusty1s / pytorch_cluster

PyTorch Extension Library of Optimized Graph Cluster Algorithms
MIT License
816 stars 146 forks source link

Improve the numerical stability of cosine similarity #134

Open chenzhekl opened 2 years ago

chenzhekl commented 2 years ago

Related: https://github.com/pytorch/pytorch/pull/31378 Current implementation: https://github.com/rusty1s/pytorch_cluster/blob/dbcafbe6a60aaa631b39050e3aa228f6d3fd1592/csrc/cuda/knn_cuda.cu#L56

As introduced in the linked PR for PyTorch, normalizing the vectors first and then doing the inner product could prevent overflow for vectors with large norms. In addition, this could ensure the calculated cosine similarity is within the range [-1, +1].

rusty1s commented 2 years ago

This is clever. Do you want to contribute this feature? :)

chenzhekl commented 2 years ago

Sure, I would love to.

rusty1s commented 2 years ago

Thanks!