Machine LearningMedium

🔍 Embeddings & Retrieval

Dense vector representations, similarity search, and retrieval systems at scale

Embeddings turn text into vectors where semantic similarity becomes geometric distance, and retrieval finds the nearest vectors to a query. Together they power search, recommendations, and RAG, and the key interview tradeoffs are bi-encoder vs cross-encoder accuracy and HNSW vs IVF-PQ index speed and memory.

What Are Embeddings?

Embeddings are dense, fixed-size vector representations that capture semantic meaning. Objects with similar meaning are mapped to nearby points in the vector space. They are the foundation of modern search, recommendation, retrieval, and clustering systems.

Types of Embeddings

Word Embeddings (legacy): Word2Vec, GloVe give one vector per word, context-independent. "bank" has the same embedding regardless of context (river bank vs. financial bank). Largely superseded by contextual embeddings.

Contextual Embeddings: BERT, GPT produce different embeddings for the same word depending on surrounding context. The [CLS] token output or mean pooling of token embeddings gives a sentence-level representation.

Sentence/Document Embeddings: Purpose-built models (Sentence-BERT, E5, BGE, Cohere Embed, OpenAI text-embedding-3) that output a single vector for an entire text passage. Trained with contrastive objectives to place semantically similar texts nearby.

Multi-modal Embeddings: CLIP, SigLIP, ImageBind map different modalities (text, image, audio) into a shared space, enabling cross-modal retrieval (search images with text queries).

How Embedding Models Are Trained

Contrastive Learning: The dominant paradigm. Given anchor-positive pairs (semantically similar texts), train the model to place them close in embedding space while pushing apart anchor-negative pairs.

  • In-batch negatives: Other items in the same mini-batch serve as negatives. Simple but can have false negatives (another item in the batch might actually be relevant).
  • Hard negatives: Mine difficult negatives, items that are similar but not relevant (e.g., documents about the same topic but answering a different question). Hard negatives dramatically improve retrieval quality.
  • InfoNCE loss: L = -log(exp(sim(q, d+)) / Σ exp(sim(q, di))), a softmax over similarities, pushes positive score above all negatives.

Matryoshka Representation Learning (MRL): Train the model so that the first k dimensions of the embedding are also a valid (lower-quality) embedding. This allows flexible dimensionality at inference time: use 256-dim for fast retrieval, 1024-dim for reranking, all from the same model. OpenAI's text-embedding-3 and Nomic Embed use MRL.

Late Interaction (ColBERT): Instead of compressing the entire document into one vector, keep per-token embeddings and compute relevance via MaxSim: the sum of maximum similarities between each query token and all document tokens. More expressive than single-vector retrieval but requires storing per-token embeddings.

Similarity Metrics

Metric Formula When to Use
Cosine similarity dot(a,b) / (‖a‖·‖b‖) Default for normalized embeddings
Dot product dot(a,b) When magnitude matters (popularity)
Euclidean (L2) ‖a-b‖ Rarely used for text; common in vision

Most embedding models normalize their outputs, making cosine similarity equivalent to dot product.

Exact nearest neighbor search is O(n), infeasible at scale. ANN algorithms trade a small amount of recall for dramatically faster search.

HNSW (Hierarchical Navigable Small World): Builds a multi-layer graph where each node connects to nearby neighbors. Search navigates from coarse (top layers) to fine (bottom layers). Best general-purpose ANN: fast search, good recall, but memory-intensive (stores the full graph in RAM).

IVF (Inverted File Index): Clusters vectors into partitions using k-means. At query time, only search the nearest k clusters. Faster indexing than HNSW, lower memory, but lower recall. Often combined with quantization (IVF-PQ).

Product Quantization (PQ): Compresses vectors by splitting them into sub-vectors and quantizing each to the nearest centroid in a codebook. Reduces memory by 10-50× with moderate recall loss. Combined with IVF: IVF-PQ is the go-to for billion-scale search.

ScaNN (Google): Anisotropic vector quantization quantizes vectors with direction-dependent error, preserving the most important dimensions. State-of-the-art recall-latency tradeoff.

Vector Databases

Purpose-built databases for storing and querying embedding vectors with metadata filtering.

Database Type Key Feature
Pinecone Managed SaaS Zero-ops, auto-scaling
Weaviate Open-source Hybrid search (dense + BM25) built-in
Qdrant Open-source Rich filtering, Rust performance
Milvus Open-source Billion-scale, GPU acceleration
pgvector Postgres extension Use existing Postgres, simple setup
Vespa Open-source Hybrid search + ML ranking in one system
Chroma Open-source Simple API, great for prototyping

Choosing a vector DB: For prototyping, use Chroma or pgvector. For production with <10M vectors, Qdrant or Weaviate. For billion-scale, Milvus or Vespa. For zero-ops SaaS, Pinecone.

Retrieval System Architecture

Two-stage retrieval is the standard production pattern:

Stage 1, Candidate Retrieval (recall-oriented):
Fast ANN search returns top-100 to top-1000 candidates. Uses bi-encoder embeddings. Latency: 5-50ms. Optimized for recall, so it's fine to return some irrelevant results.

Stage 2, Reranking (precision-oriented):
A cross-encoder (e.g., Cohere Rerank, BGE-reranker, ColBERT) scores each query-candidate pair with full attention. Reduces top-1000 to top-10. Latency: 50-300ms. Cross-encoders are much more accurate than bi-encoders but can't be pre-indexed.

Why two stages? Bi-encoders can pre-compute document embeddings (index once, query many times) but use independent encoding, with no query-document interaction until the dot product. Cross-encoders jointly encode query + document with full cross-attention: far more accurate, but they must run at query time for each candidate, so they can only process a small set.

Combining dense (embedding-based) retrieval with sparse (BM25/keyword) retrieval. Dense search excels at semantic matching and paraphrases but struggles with exact terms, IDs, and rare words. Sparse search handles exact matching well.

Reciprocal Rank Fusion (RRF): Merge two ranked lists without needing score calibration: score(d) = Σ 1/(k + rank_i(d)) where k is typically 60. Simple, robust, widely used.

Evaluation Metrics

  • Recall@k: Fraction of relevant documents in the top-k results
  • MRR (Mean Reciprocal Rank): Average of 1/rank of the first relevant result
  • NDCG@k: Accounts for graded relevance, not just binary
  • Hit Rate: Is the relevant document anywhere in top-k? (binary recall)
  • Latency: p50/p99 query time, critical for user-facing search

Scaling Embeddings to Billions

At billion-scale, challenges multiply:

  • Index size: 1B × 1024-dim × 4 bytes = 4TB in RAM for HNSW. Solution: IVF-PQ reduces to ~100-200GB.
  • Index build time: HNSW over 1B vectors takes days. Solution: distributed indexing with sharding.
  • Freshness: New documents need to be indexed quickly. Solution: tiered indexes, a small "hot" HNSW index for recent docs merged periodically into the main IVF-PQ index.
  • Filtering: "Find similar documents, but only from 2024, in English, by verified authors." Solution: pre-filtering (filter before ANN, reduces recall) or post-filtering (filter after ANN, may return fewer than k results). Qdrant and Vespa support efficient filtered search.