Implement Vector Similarity Search

Math Foundations

Cosine Similarity measures the cosine of the angle $\theta$ between Query vector $q \in \mathbb{R}^D$ and Document vector $v \in \mathbb{R}^D$:

$$\text{CosineSimilarity}(q, v) = \cos(\theta) = \frac{q \cdot v}{| q |_2 | v |2} = \frac{\sum{j=1}^D q_j v_j}{\sqrt{\sum q_j^2} \sqrt{\sum v_j^2}}$$

• Scores range between $-1.0$ (Opposite directions) and $+1.0$ (Identical directions).
• Scale Invariance: Cosine similarity measures angular orientation, independent of vector magnitude length.

L2 Normalization Speed Trick

If we pre-normalize all database vectors $v_i$ and query vector $q$ to Unit L2 Length ($| v |_2 = 1, | q |_2 = 1$):

$$\hat{v} = \frac{v}{| v |_2}, \quad \hat{q} = \frac{q}{| q |_2}$$

The denominator simplifies to $1 \cdot 1 = 1$, and Cosine Similarity reduces to a simple Dot Product:

$$\text{CosineSimilarity}(\hat{q}, \hat{v}) = \hat{q} \cdot \hat{v}$$

We can query $N$ database vectors $V \in \mathbb{R}^{N \times D}$ in parallel using a single Matrix-Vector Product:

$$\text{Scores} = V_{\text{norm}} \hat{q} \in \mathbb{R}^N$$

NumPy Implementation from Scratch

import numpy as np


class VectorSimilaritySearchFromScratch:

  def __init__(self, use_normalization=True):
    self.use_normalization = use_normalization
    self.vectors_norm = None
    self.doc_ids = None

  def _l2_normalize(self, matrix):
    # L2 normalize rows along feature axis 1
    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
    # Avoid division by zero
    norms = np.maximum(norms, 1e-12)
    return matrix / norms

  def add_documents(self, doc_ids, vectors):
    # vectors shape: [N, D]
    self.doc_ids = np.array(doc_ids)

    if self.use_normalization:
      self.vectors_norm = self._l2_normalize(vectors)
    else:
      self.vectors_norm = vectors

  def search(self, query_vector, k=5):
    # query_vector shape: [D]
    query_vector = np.array(query_vector).reshape(1, -1)

    if self.use_normalization:
      query_norm = self._l2_normalize(query_vector)
    else:
      query_norm = query_vector

    # Compute Cosine Similarities across all N documents in parallel
    # Similarity scores shape: [N]
    scores = np.dot(self.vectors_norm, query_norm.T).flatten()

    # Find Top K highest similarity indices using argpartition in O(N)
    if len(scores) <= k:
      top_k_indices = np.argsort(-scores)
    else:
      top_k_indices = np.argpartition(-scores, k)[:k]
      # Sort top k indices in exact descending order
      top_k_indices = top_k_indices[np.argsort(-scores[top_k_indices])]

    top_ids = self.doc_ids[top_k_indices]
    top_scores = scores[top_k_indices]

    return list(zip(top_ids, top_scores))

Say this out loud

Vector similarity search ranks document embeddings against query vectors using Cosine Similarity. Pre normalizing vectors to unit L2 length simplifies Cosine Similarity to high speed matrix dot products. Matrix vector multiplication scores N database vectors in parallel, and top K partitioning extracts top matches efficiently.

Followups to expect

1. What is Euclidean Distance vs Cosine Similarity on L2 normalized vectors? On unit L2 normalized vectors, Euclidean Distance and Cosine Similarity are monotonic inverses: $| \hat{q} - \hat{v} |^2 = 2 - 2 (\hat{q} \cdot \hat{v})$.
2. When does exact brute-force search become too slow? Brute force matrix search takes $\mathcal{O}(N \cdot D)$ time. When catalog size $N > 1,000,000$, Approximate Nearest Neighbor (ANN) indexes like HNSW are required for sub-millisecond search.