Dense Retrieval and Recall@k

The problem

In dense retrieval, a neural encoder turns every passage, and the query, into a vector of numbers (an embedding). Passages whose vectors point in the same direction as the query's are about the same thing. You will rank passages by the cosine similarity of their embeddings with the query's, and measure how good the ranking is.

Cosine similarity

The cosine similarity of two vectors is their dot product after scaling each to unit length. It is 1 if they point the same way, 0 if they are unrelated (perpendicular), and $-1$ if opposite. Only the direction counts, never the length.

cosine_scores(queries, passages) returns the table with one row per query and one column per passage.

Example. Two queries $(1, 0)$ and $(0, 1)$, and four passages $(2, 0)$, $(1, 1)$, $(0, 3)$ and $(-1, 0.1)$. Cosine similarities:

• query 0 — passage 0: $1.00$; passage 1: $0.71$; passage 2: $0.00$; passage 3: $-0.995$
• query 1 — passage 0: $0.00$; passage 1: $0.71$; passage 2: $1.00$; passage 3: $0.10$

Passage 0, $(2, 0)$, is twice as long as the query but points the same way: similarity exactly 1.

The top $k$

top_k(scores, k) gives, for every query, the indices of the k highest scores, best first. On ties, the lower index comes first.

Example. For the table above and $k = 2$: query 0 gets [0, 1], query 1 gets [2, 1].

Recall at $k$

The recall@k of a query is the fraction of its relevant passages that appear in its top $k$. recall_at_k(retrieved, relevant) takes the table from top_k and a list of sets of relevant passage indices, and returns the mean over queries. A query with no relevant passage is skipped.

Example. Relevant passages: $\{0, 1\}$ for query 0 and $\{2, 3\}$ for query 1. Query 0 retrieved [0, 1]: both found, recall $1$. Query 1 retrieved [2, 1]: passage 2 found, 3 missed, recall $0.5$. Mean: $0.75$.

What to implement

• cosine_scores(queries, passages): queries has shape $(Q, D)$ and passages $(P, D)$. Return the $(Q, P)$ table of cosine similarities.
• top_k(scores, k): shape $(Q, k)$.
• recall_at_k(retrieved, relevant): the mean recall.

Requirements

• No loops over passages. A loop over queries is fine in recall_at_k.

Hint

np.argsort(-scores, axis=1, kind="stable") sorts each row from the highest score down, with ties in index order. To scale rows to unit length, divide by np.linalg.norm(x, axis=-1, keepdims=True).