Cosine Similarity Top-K Retrieval

1. Background Knowledge

Cosine Similarity measures the cosine of the angle between two non-zero vectors in an inner product space. It is a measurement of orientation and not magnitude, making it ideal for comparing high-dimensional embeddings where direction matters more than length. In Vision-Language Models (VLMs) like CLIP, text and image embeddings are projected into a shared latent space. The contrastive learning objective used to train these models aligns matching pairs by maximizing their cosine similarity, effectively normalizing the vectors during training or inference. Consequently, the magnitude of the embedding vectors is often arbitrary or normalized to unit length, rendering raw dot products unreliable if vectors are not pre-normalized.

The formula for cosine similarity between a query vector $q$ and a gallery vector $g_i$ is:

where $\lVert v \rVert_2 = \sqrt{\sum_{j=1}^{d} v_j^2}$ is the L2 norm (Euclidean length). If both vectors are already unit-normalized ($\lVert q \rVert_2 = 1$ and $\lVert g_i \rVert_2 = 1$), the denominator becomes 1, and cosine similarity simplifies to the dot product $q \cdot g_i$. However, in general retrieval tasks, you must explicitly compute the norms to ensure correctness unless the problem guarantees pre-normalized inputs.

Top-K Retrieval is a fundamental operation in information retrieval and recommendation systems. It involves computing a relevance score for every candidate in a database (gallery) and returning the $k$ candidates with the highest scores. This is essentially a nearest-neighbor search problem. When $k$ is small relative to the gallery size $m$, efficient algorithms can avoid sorting the entire list, but for moderate $m$, a full sort or a partial sort is often sufficient and simpler to implement.

2. Algorithm Approach

The core approach is a brute-force nearest-neighbor search using cosine similarity. Since the gallery size $m$ is typically manageable in "Easy" level problems, we can compute the similarity score for every candidate in the gallery against the query.

1. Normalization: Compute the L2 norm of the query vector and each gallery vector.
2. Scoring: For each gallery vector, compute the dot product with the query, then divide by the product of their norms to get the cosine similarity.
3. Ranking: Store the similarity scores along with their original indices.
4. Selection: Sort the candidates based on their similarity scores in descending order.
5. Tie-Breaking: Ensure that if two candidates have the same similarity score, the one with the lower original index appears first. This is a stable sort requirement.
6. Extraction: Take the top $k$ entries from the sorted list and return their indices and scores.

3. Step-by-Step Strategy

1. Compute Query Norm: Calculate the L2 norm of the query vector. Let this be $n_q$.
2. Initialize Results List: Create a list to store tuples of (similarity_score, original_index).
3. Iterate Through Gallery:

• For each vector $g_i$ in gallery at index $i$:
• Compute the L2 norm of $g_i$, let this be $n_i$.
• Compute the dot product: $\text{dot} = \sum_{j=0}^{d-1} q[j] \times g_i[j]$.
• Compute cosine similarity: $\text{sim} = \frac{\text{dot}}{n_q \times n_i}$.
• Append (sim, i) to the results list.

4. Sort the Results:

• Sort the list of tuples.
• Primary key: Similarity score in descending order.
• Secondary key: Original index in ascending order (to handle ties correctly).
• In Python, you can achieve this by sorting with a key that negates the score for descending order, or by using reverse=True and ensuring the sort is stable (Python's sort is stable). If using reverse=True, you should sort by (-score, index) or simply rely on stability if you sort by score descending and the original list was ordered by index. A robust way is to sort by (-sim, i).

5. Extract Top-K:

• Slice the first $k$ elements from the sorted list.
• Separate the indices and scores into two lists.

6. Format Output:

• Round each similarity score to 4 decimal places.
• Return [indices, sims].

4. Common Pitfalls

• Division by Zero: If a gallery vector is a zero vector (all zeros), its norm is 0, leading to a division by zero error. While rare in valid embedding spaces, it's good practice to handle or assume non-zero vectors.
• Incorrect Tie-Breaking: The problem specifies that ties must break toward the lower index. If you simply sort by score descending, the relative order of tied elements depends on the sort algorithm's stability. Python's sort is stable, meaning if you sort by score descending, elements with equal scores will retain their original relative order (which is ascending index if you iterated in order). However, explicitly including the index in the sort key (e.g., key=lambda x: (-x, x)) is safer and clearer.
• Floating Point Precision: Cosine similarity calculations involve floating-point arithmetic. Small errors can occur. Ensure you round the final output to 4 decimal places as specified, but do not round intermediate calculations prematurely.
• Confusing Dot Product with Cosine Similarity: If the vectors are not unit-normalized, the dot product is not the cosine similarity. Always divide by the product of the norms unless you are certain the inputs are pre-normalized.
• Index Off-by-One: Ensure you are returning the original indices from the gallery list, not the indices from the sorted list.

5. Time & Space Complexity

• Time Complexity:
• Computing norms and dot products for all $m$ gallery vectors takes $O(m \cdot d)$, where $d$ is the dimensionality of the vectors.
• Sorting the $m$ results takes $O(m \log m)$.
• Total time complexity is $O(m \cdot d + m \log m)$. For typical embedding dimensions ($d \approx 512$ or $768$) and moderate gallery sizes, this is efficient.
• Space Complexity:
• We store a list of $m$ tuples (score, index), which takes $O(m)$ space.
• The output lists take $O(k)$ space.
• Total space complexity is $O(m)$.