O(N log N) algorithm/implementation for maximum inner product search

Question

I am in a situation where I have a query space Q and a key space K, both filled with N d-dimensional real vectors (N ≈ 10^6, d ≈ 50). For each query q, I want to find the k≈10 keys k_i that have the highest inner products with q. One may assume that both keys and queries are normally distributed.

It seems that getting the k nearest neighbours for one query should be possible in O(log N) time by creating some kind of binary search tree on the keys, which would lead to an O(N log N) time in total. And if query computations can be combined in some way, this could lead to a further speedup.

However, I have found no algorithm or implementation yet that advertises itself as being O(N log N). If there exists such an algorithm, could you provide a reference? If there does not, could you provide a reason(ing) for that?

Answer 1

For your parameters, I suspect it will be hard to do better than a linear scan, i.e., $O(Nd)$ time per query vector.

This problem is about as hard the problem of nearest neighbor search.

Specifically, when $A,B$ are two vectors that have length 1, their inner product satisfies

$$\langle A, B \rangle = 1 - \|A-B\|_2^2/2.$$

(See https://math.stackexchange.com/q/1236465/14578.) Therefore, in this case, the pairs with highest inner product are the pairs with lowest distance.

Unfortunately, finding nearest neighbors in high dimension is often hard. When $d \approx 50$, I suspect that it might be hard to beat a linear scan. I'm not sure that fancy data structures will help.

Since the problem is hard when computing smallest distances, by the above relation, I expect finding largest inner products will also be hard. In particular, this reduction implies that finding largest inner products in the special case where all vectors have length 1 is likely to be hard (it's exactly as hard as finding nearest neighbors). If your problem is hard in a special case, it is also hard in the general case (assuming you care about worst-case running time).

See also K Nearest Neighbor using K-D Tree (What ratio makes K-D beneficial)?, How can I make k nearest neighbor queries fast on unit hypersphere?, How do I choose an optimal cell size when searching for close pairs of points, and using cells to implement this? for more about nearest neighbor search. There are many fancy data structures that have been devised for nearest neighbor search. I don't know whether they can be generalized to apply to search for highest inner product.

Answer 2

I don't think any exact approach to solve this problem in $O(N \log N)$ exists. You can look for Approximate Nearest Neighbors search methods with theoretical logarithmic complexity for a single query (e.g. HNSW, NSG, etc). However these guarantees do not hold in practice because you lack the hypothesis on which the proofs are based. In practice you can regulate the trade-off between accuracy and speed of the search through some parameters (e.g. ef_search in HNSW).

These approaches are very fast and effective in practice and, by tuning the search parameters you can easily achieve accuracies near $100\%$. They are really fast because they can find the approximate $k$NN by scanning a small percentage of the dataset (e.g. $<1\%$). Also, they allow to submit many queries at the same time and they benefit from handling multiple queries thanks to techniques like for instance vectorization.

I suggest you to try FAISS (Facebook AI Similarity Search) which has many methods for ANN search (also methods like IVF index or hash-based approaches or GPU indexes that are very fast on GPUs).

Note that the guarantees hold for distance metrics, while inner product is not a distance. However, very often it also works in practice and it is implemented as one possible metric to compare vectors.