Cosine Top-k Embedding Search - Representation Learning & Retrieval

Dense Retrieval with Cosine Similarity

What you'll build

Exact top-k nearest-neighbour search by cosine similarity: for each query embedding, return the k most similar documents and their scores.

Background

This is the "retrieve" step of a RAG system and of every semantic search engine. Embeddings come from a model trained with a contrastive loss (previous problem). With L2-normalized vectors, cosine similarity is a plain dot product, so scoring all documents for a batch of queries is one matrix multiply. Vector databases use approximate indexes (HNSW graphs, IVF clustering, product quantization) to avoid scoring every document; this brute-force version is what they approximate, and what their recall is measured against.

Algorithm

q_n    = queries / max(||queries||, 1e-12)     # (Q, D), per row
d_n    = docs / max(||docs||, 1e-12)           # (N, D), per row
scores = q_n @ d_n.T                           # (Q, N) cosine similarities

For each query row, sort the documents by score descending and keep the first k.

Inputs

• queries: float array (Q, D)
• docs: float array (N, D)
• k: int, 1 <= k <= N

Expected Return Format

A tuple (indices, scores):

• indices: (Q, k) integer array of document indices, best first
• scores: (Q, k) array of the matching cosine similarities, non-increasing along each row

Edge cases and conventions

• Ties: equal scores keep the lower document index first. A stable sort of the negated scores does this; np.argpartition and torch.topk don't guarantee it.
• Documents pointing the same way score the same regardless of their length (that's the point of cosine).

Constraints

NumPy only, vectorised over queries and documents.

Follow-up

• Use np.argpartition to get the top-k in O(N) per query, then sort only those k. How do you keep the tie rule?
• When is maximum inner product (no normalization) the right similarity instead of cosine?

Input: cosine_topk(np.array([[1.0, 0.0]]), np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [-1.0, 0.0]]), k=2)
Output: (indices = [[0, 2]], scores = [[1.0, 0.7071]])

Input: cosine_topk(np.array([[1.0, 0.0]]), np.array([[0.0, 3.0], [2.0, 0.0], [1.0, 1.0], [1.0, 0.0], [5.0, 0.0]]), k=4)[0]
Output: [[1, 3, 4, 2]]

import numpy as np

def cosine_topk(queries, docs, k):
    """
    Exact top-k retrieval by cosine similarity.

    Args:
        queries: (Q, D) query embeddings
        docs: (N, D) document embeddings
        k: number of results per query

    Returns:
        (indices (Q, k), scores (Q, k)), sorted by score descending,
        ties broken by lower document index
    """
    # Your code here
    pass