# Fastest way to do Maximum Inner Product Search?

There have lots of $$d-$$dimension vectors, they combine a set $$X$$. We use an input query $$q$$. We need to find the $$p$$, which has the maximum inner product from the set $$X$$.

$$p = \arg \max_{x \in X} x^T \ q$$

import torch
import time

def main():

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

matrix_size = 800000000
dim_size = 768
batch_size = 100000  # Adjust the batch size based on available memory

# create a query vector using torch
query = torch.randn(1, dim_size, dtype=torch.float32).to(device)

print("Query vector:", query)

# start timing
start = time.time()

for i in range(0, matrix_size, batch_size):
# create a batch of the matrix
matrix_batch = torch.randn(batch_size, dim_size, dtype=torch.float32).to(device)

# calculate the inner product for the batch
result_batch = torch.matmul(matrix_batch, query.T)

# torch topk, k = 10 for the batch
values_batch, indices_batch = torch.topk(result_batch.squeeze(), k=10)

# process the top-k results as needed
print("Top values in batch:", values_batch)
print("Indices in batch:", indices_batch)

# end timing
end = time.time()
print("Time elapsed: ", end - start)

if __name__ == "__main__":
main()


Similar with nearest neighbor question, changing the definition to find a $$p$$, which makes the distance between $$q$$ to become closer.

Here the distance set as Euclidean distance.

$$p = \arg \min_{x \in X} \| q - x \|^2 = (\| x \|^2 - 2q^T x)$$

If the vector module in $$X$$ are all the same, then the two problems are actually equivalent.

However, in many practical scenarios, such as BERT-encoded sentence vectors, various Embeddings in recommendation systems, etc., this constraint is not satisfied.

#### Locality-sensitive hashing

LSH seems to be a big step forward compared to the naive algorithm. But don’t be too happy. To achieve the effect of $$O(d \log N)$$ ($$N$$ is the # documents), you must obey that strong assumption. It is often unrealistic for points to be distributed evenly in space. In addition to this, an LSH can only be suitable for certain distance measures. For MIPS, no LSH that meets the requirements can be found.

#### Asymmetric LSH(ALSH)

Can anyone suggest an algorithm? Thanks

• Have you considered using a vector database like faiss? They are meant to address the problem you described. However, as you use PyTorch in your code, I am afraid that you may have other requirements that didn't mention explicitly in your question.
– noe
Nov 16, 2023 at 9:29
• @noe We just want to search relative document based on the query. I think faiss may help us. Thanks. Nov 16, 2023 at 13:19