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__":

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

  • $\begingroup$ 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. $\endgroup$
    – noe
    Nov 16, 2023 at 9:29
  • $\begingroup$ @noe We just want to search relative document based on the query. I think faiss may help us. Thanks. $\endgroup$
    – jackson
    Nov 16, 2023 at 13:19
  • $\begingroup$ Great! I added an answer with more information on the matter so that the question can be marked as resolved. $\endgroup$
    – noe
    Nov 16, 2023 at 15:00

1 Answer 1


If what you want is a fast way of vector-search documents, then a vector database may fit your needs.

Some popular ones include faiss and pinecone.

If you already have a PostgreSQL database, you may also consider using the pgvector extension.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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