Representation Learning, Embeddings, and Similarity

Embeddings are the lingua franca of modern ML systems. Every major product feature that "understands" content - semantic search, recommendation, deduplication, clustering, anomaly detection - runs on the same primitive: a dense vector that encodes meaning as position in a geometric space. Items that mean similar things sit close together. Items that mean different things sit far apart. Distance is semantic similarity.

This module covers how embeddings are learned, how similarity search works at scale, and how to wire both into production retrieval systems.

What Embeddings Actually Are

An embedding model is a neural network whose final layer outputs a fixed-size vector rather than class probabilities. The model is trained so that the geometry of the vector space encodes semantic relationships.

The two dominant training approaches:

Contrastive learning: Push representations of similar items together, dissimilar items apart. Used for image encoders (SimCLR, MoCo) and sentence encoders.

Supervised metric learning: Use labeled pairs (query, positive document) and train with a loss that measures inter-class distances. Used for search and recommendation.

pythonCopy
import torch
import torch.nn as nn
import torch.nn.functional as F

class EmbeddingModel(nn.Module):
    def __init__(self, input_dim: int, embed_dim: int = 128):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.GELU(),
            nn.Linear(256, embed_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        emb = self.encoder(x)
        return F.normalize(emb, p=2, dim=-1)  # L2-normalize to unit sphere

L2 normalization is almost always applied to the output: it constrains all embeddings to the unit hypersphere, making cosine similarity equivalent to dot product. This simplifies search and loss computation.

Triplet Loss and Contrastive Loss

Triplet loss trains the model with triples (anchor, positive, negative). It pushes the anchor closer to the positive than to the negative by at least a margin $m$:

$\mathcal{L} = \max(d(a, p) - d(a, n) + m, 0)$

pythonCopy
class TripletLoss(nn.Module):
    def __init__(self, margin: float = 0.3):
        super().__init__()
        self.margin = margin

    def forward(self, anchor, positive, negative):
        d_ap = F.pairwise_distance(anchor, positive)
        d_an = F.pairwise_distance(anchor, negative)
        loss = F.relu(d_ap - d_an + self.margin)
        return loss.mean()

# Mining hard negatives: pick negatives that are close to the anchor
# (semi-hard negatives) rather than random - dramatically improves convergence
def mine_hard_negatives(anchor_emb, all_emb, k=5):
    sims = anchor_emb @ all_emb.T      # cosine sims (both L2-normalized)
    # Exclude anchor itself, pick top-k most similar non-positive
    _, indices = sims.topk(k + 1, dim=-1)
    return indices[:, 1:]              # drop self

Hard negative mining is essential for embedding quality. Random negatives are too easy - the model can separate them without learning meaningful structure. Hard negatives force the model to learn fine-grained distinctions.

Approximate Nearest Neighbor Search with FAISS

For retrieval at scale, exact nearest-neighbor search (compare query against every vector) is too slow. FAISS provides approximate nearest neighbor (ANN) indexes that trade small accuracy loss for orders-of-magnitude faster search:

pythonCopy
import faiss
import numpy as np

def build_faiss_index(embeddings: np.ndarray, use_gpu: bool = False):
    d = embeddings.shape[1]          # embedding dimension

    # IVF (Inverted File) index: cluster the space into n_list Voronoi cells,
    # search only the nprobe nearest cells at query time
    n_list = min(4096, int(np.sqrt(len(embeddings))))
    quantizer = faiss.IndexFlatIP(d)  # Inner product (cosine on L2-normalized)
    index = faiss.IndexIVFFlat(quantizer, d, n_list, faiss.METRIC_INNER_PRODUCT)

    index.train(embeddings.astype(np.float32))
    index.add(embeddings.astype(np.float32))
    index.nprobe = 64                 # Search 64 cells; increase for more recall

    if use_gpu:
        res = faiss.StandardGpuResources()
        index = faiss.index_cpu_to_gpu(res, 0, index)

    return index

def search(index, query_emb: np.ndarray, k: int = 10):
    q = query_emb.reshape(1, -1).astype(np.float32)
    scores, ids = index.search(q, k)
    return ids[0], scores[0]

FAISS index types to know:

• IndexFlatIP: exact search, no training. Baseline for correctness checks.
• IndexIVFFlat: cluster-based ANN. Good to ~10M vectors.
• IndexIVFPQ: adds product quantization compression. Good for 100M+ vectors where RAM is constrained.
• IndexHNSW: graph-based ANN. Very fast queries; higher build cost.

Sentence Embeddings with Sentence Transformers

For text, use pretrained sentence embeddings rather than training from scratch:

pythonCopy
from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("BAAI/bge-small-en-v1.5")  # 33M params, fast, strong

texts = [
    "How do I cancel my subscription?",
    "I want to stop my monthly plan",
    "What is the refund policy?",
]

embeddings = model.encode(texts, normalize_embeddings=True, batch_size=64)
# embeddings.shape: (3, 384)

# Cosine similarity matrix
sims = embeddings @ embeddings.T
print(sims)
# [[1.00, 0.87, 0.43],  ← "cancel" ↔ "stop plan" are close; "refund" is further
#  [0.87, 1.00, 0.41],
#  [0.43, 0.41, 1.00]]

Good pretrained embedding models for production (ranked by size/speed):

• bge-small-en-v1.5 (33M): fast, strong for retrieval
• bge-base-en-v1.5 (109M): balanced
• text-embedding-3-small (OpenAI API): easy to use, high quality, API cost
• all-mpnet-base-v2 (110M): general purpose, good off-the-shelf

pgvector for Production Retrieval

FAISS is an in-memory library. For a production service with a PostgreSQL database, use pgvector:

sqlCopy
-- Enable extension
CREATE EXTENSION IF NOT EXISTS vector;

-- Store embeddings alongside metadata
CREATE TABLE documents (
    id         BIGSERIAL PRIMARY KEY,
    content    TEXT,
    embedding  vector(384),
    created_at TIMESTAMPTZ DEFAULT now()
);

-- Build an index for approximate search
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- Nearest-neighbor retrieval
SELECT id, content, 1 - (embedding <=> $1::vector) AS similarity
FROM documents
ORDER BY embedding <=> $1::vector
LIMIT 10;

pgvector eliminates the need for a separate vector database service. For most applications under a few million documents, it is the right operational choice: one database, one backup strategy, SQL joins with your existing tables.

Common Mistakes and Bad Instincts

Not L2-normalizing embeddings before cosine similarity. Raw embeddings have variable norms. Cosine similarity between un-normalized vectors conflates magnitude with direction. Normalize to the unit sphere before storing or computing similarity.

Training embeddings from scratch when pretrained ones exist. For text, images, and many structured domains, strong pretrained embeddings are freely available. Attempting to train from scratch on small datasets almost always produces inferior representations.

Choosing FAISS over pgvector reflexively. FAISS is an excellent library, but it requires a separate service, a separate operational burden, and loses SQL expressiveness. For most products, pgvector in Postgres is the right first choice.

Not tuning nprobe in FAISS IVF indexes. Default nprobe = 1 produces fast but low-recall search. Run recall benchmarks against ground-truth exact search and tune nprobe to meet your target. A typical production setting is 32–128.

Where to Go Next

• retrieval-systems-vector-databases-and-rag: combine embeddings with retrieval pipelines to build RAG systems
• transformers-and-modern-nlp-for-engineers: understand the Transformer architecture that produces the embeddings you are using
• fine-tuning-adaptation-and-when-not-to-fine-tune: fine-tune an embedding model on domain-specific data when off-the-shelf quality is insufficient