Modern Embedding Architectures & Selection

You shall know a word by the company it keeps.

Vec, Linguistically Social AI Agent

31.2.1 Modern Embedding Architectures

Matryoshka Representation Learning (MRL)

Traditional embedding models produce fixed-dimension vectors (e.g., 768 or 1536 dimensions). If you need smaller vectors for storage efficiency, you must train a separate model. Matryoshka Representation Learning (Kusupati et al., 2022) solves this by training a single model whose embeddings are useful at multiple dimensionalities. The first d dimensions of a 768-dimensional embedding form a valid d-dimensional embedding, much like nested Russian dolls.

During training, the loss function is computed at multiple truncation points (e.g., dimensions 32, 64, 128, 256, 512, 768), and the gradients from all truncation levels are summed. This forces the model to pack the most important information into the leading dimensions. Figure 31.2.1 visualizes this nested structure.

Algorithm: MRL contrastive training
Input:  encoder f_theta producing full-dim embeddings z = f_theta(x) in R^d
        set of nested dimensionalities M = {d_1 < d_2 < ... < d_K = d}
        loss family L(.) (e.g., InfoNCE for retrieval, cross-entropy for classification)
        per-level weights w_1, .., w_K (typically uniform, sum to 1)
Output: encoder whose first d_m coordinates form a valid d_m-embedding for every m

  // For every minibatch (x, x_pos, {x_neg^j})
  z      := f_theta(x)
  z_pos  := f_theta(x_pos)
  z_neg  := f_theta(x_neg^j) for j

  L_total := 0
  For m = 1..K:
      // Truncate all three vectors to the first d_m coordinates,
      // re-L2-normalize the truncated vector.
      z_m     := normalize(z[:d_m])
      z_pos_m := normalize(z_pos[:d_m])
      z_neg_m := normalize(z_neg[:d_m])

      L_m := L( z_m, z_pos_m, z_neg_m )                // e.g. InfoNCE
      L_total := L_total + w_m * L_m

  Backprop on L_total; update theta

Effect. Summing the truncated losses forces the network to put the most discriminative
content in the first d_1 coords, the next-most-discriminative in d_1..d_2, and so on.
At inference any d_m can be used by simple slicing + renormalization; no separate
projection head, no re-training.

Practical menu (OpenAI text-embedding-3, Nomic v1.5, BGE-M3):
   d in {64, 128, 256, 512, 768 (or 1024, 1536, 3072)}
   storage savings: d_m / d_K, e.g. 64 / 1536 ~= 4.2 percent of full storage,
   retrieval recall at d_m = 64 is typically 85..92 percent of full-dim recall.

Figure 31.2.1b: Matryoshka embeddings pack information hierarchically, enabling flexible dimensionality reduction by truncation.

ColBERT: Late Interaction

ColBERT introduces a late interaction paradigm that sits between bi-encoders and cross-encoders. Instead of compressing each sentence into a single vector, ColBERT retains per-token embeddings for both the query and the document. At search time, it computes the maximum similarity between each query token and all document tokens (MaxSim), then sums these scores across query tokens.

Formally, given query token embeddings $\{q_1, \dots, q_m\}$ and document token embeddings $\{d_1, \dots, d_n\}$ (both L2-normalised), the ColBERT score is

that is, for every query token pick its most similar document token and sum those maxima. This decomposition is what makes ColBERT a "late" interaction: the per-token embeddings are computed independently (so document side can be precomputed), but the interaction between query and document still happens at the token level, not the sentence level.

This approach preserves much of the expressiveness of cross-encoders while still allowing document token embeddings to be precomputed. The tradeoff is storage: a document with 200 tokens requires storing 200 vectors instead of one. ColBERT v2 addresses this through residual compression, reducing storage by 6 to 10 times while preserving retrieval quality.

# ColBERT-style MaxSim scoring (simplified)
import torch
import torch.nn.functional as F
def colbert_score(query_tokens, doc_tokens):
    """
    Compute ColBERT late-interaction score.
    query_tokens: (num_query_tokens, dim) - per-token query embeddings
    doc_tokens: (num_doc_tokens, dim) - per-token document embeddings
    """
    # Normalize token embeddings
    query_tokens = F.normalize(query_tokens, dim=-1)
    doc_tokens = F.normalize(doc_tokens, dim=-1)
    # Similarity matrix: (num_query_tokens, num_doc_tokens)
    sim_matrix = torch.matmul(query_tokens, doc_tokens.T)
    # MaxSim: for each query token, find max similarity to any doc token
    max_sim_per_query_token = sim_matrix.max(dim=-1).values
    # Sum across query tokens
    score = max_sim_per_query_token.sum()
    return score
# Example
num_q_tokens, num_d_tokens, dim = 8, 50, 128
q_embs = torch.randn(num_q_tokens, dim)
d_embs = torch.randn(num_d_tokens, dim)
score = colbert_score(q_embs, d_embs)
print(f"ColBERT score: {score.item():.4f}")

31.2.2 Embedding Model Ecosystem and Selection

You now understand how embeddings are produced; the next decision is which embedding model to actually use. The market splits cleanly into two camps: managed API services (OpenAI, Cohere, Voyage) that trade per-call cost for zero infrastructure, and self-hosted open-weight models (E5, BGE, GTE) that trade upfront engineering for predictable cost at scale. We treat the two paths in turn.

API Embedding Services

For teams that prefer managed solutions, several providers offer embedding APIs. OpenAI's text-embedding-3-small and text-embedding-3-large models support Matryoshka-style dimension reduction through a dimensions parameter. Cohere's embed-v3 models support separate input types for queries and documents, which can improve retrieval quality. Google's Vertex AI provides the Gecko embedding model with built-in task-type specification.

# Using OpenAI embeddings with dimension control
from openai import OpenAI
import numpy as np
client = OpenAI()
texts = [
    "Retrieval augmented generation combines search with LLMs.",
    "RAG systems ground language model outputs in retrieved documents.",
    "The weather forecast predicts rain tomorrow."
]
# Generate embeddings at reduced dimensionality (Matryoshka-style)
response = client.embeddings.create(
    model="text-embedding-3-small",
    input=texts,
    dimensions=256 # Reduce from default 1536 to 256
)
embeddings = np.array([item.embedding for item in response.data])
print(f"Shape: {embeddings.shape}")
# Compute cosine similarities
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
normalized = embeddings / norms
similarity = np.dot(normalized, normalized.T)
print(f"RAG-related similarity: {similarity[0][1]:.4f}")
print(f"Unrelated similarity: {similarity[0][2]:.4f}")

The MTEB Benchmark

The Massive Text Embedding Benchmark (MTEB) evaluates embedding models across multiple tasks: classification, clustering, pair classification, reranking, retrieval, semantic textual similarity (STS), and summarization. MTEB covers dozens of datasets across multiple languages, making it the standard reference for comparing embedding models.

31.2.3 Embedding Space Geometry

Picking a model is half the work; using it correctly is the other half. Two geometric choices, which similarity metric to score against and how to handle the curse of dimensionality, can swing retrieval quality by several MTEB points without any change to the underlying model. We start with similarity metrics because they sit at every query path.

Similarity Metrics

The choice of similarity metric affects both retrieval quality and computational performance. The three standard metrics are:

• Cosine similarity: Measures the angle between vectors, ignoring magnitude. Ranges from -1 to 1. Most commonly used for text embeddings because it is invariant to vector length. Equivalent to dot product on L2-normalized vectors.
• Dot product (inner product): Measures both directional similarity and magnitude. Useful when vector length carries meaningful information (e.g., passage importance). Faster than cosine similarity when vectors are not pre-normalized.
• Euclidean distance (L2): Measures straight-line distance in the embedding space. Lower values indicate greater similarity. Equivalent to cosine distance for normalized vectors, but can behave differently for unnormalized vectors.

Dimensionality and the Curse of Dimensionality

In high-dimensional spaces, distances between points tend to concentrate: the difference between the nearest and farthest neighbor shrinks relative to the overall distance. This phenomenon, known as the curse of dimensionality, means that exact nearest neighbor search becomes less meaningful as dimensionality grows. In practice, this is why approximate nearest neighbor (ANN) algorithms work so well for embeddings: the approximation error introduced by ANN is often smaller than the noise inherent in high-dimensional distance computations. Figure 31.2.3a shows how semantic clustering emerges in embedding space.

Figure 31.2.3d: Embedding spaces organize semantically related concepts into clusters, enabling nearest-neighbor retrieval.

31.2.4 Fine-Tuning Embeddings for Domain Specificity

General-purpose embedding models may not capture domain-specific terminology or relationships well. For example, in a legal domain, "consideration" means something entirely different from its everyday usage. Fine-tuning an embedding model on domain-specific data can substantially improve retrieval quality.

# Fine-tuning a sentence transformer on domain-specific data
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    losses,
)
from datasets import Dataset
# Prepare training data: (anchor, positive) pairs
# Hard negatives are automatically mined from in-batch examples
train_data = Dataset.from_dict({
    "anchor": [
    "What is consideration in contract law?",
    "Define the doctrine of estoppel",
    "What constitutes a breach of fiduciary duty?",
    ],
    "positive": [
    "Consideration is something of value exchanged between parties to a contract.",
    "Estoppel prevents a party from asserting a claim inconsistent with prior conduct.",
    "A fiduciary breach occurs when a fiduciary acts against the beneficiary interest.",
    ],
})
# Load base model
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
# Configure training
loss = losses.MultipleNegativesRankingLoss(model)
args = SentenceTransformerTrainingArguments(
    output_dir="./legal-embedding-model",
    num_train_epochs=3,
    per_device_train_batch_size=32,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    fp16=True,
)
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_data,
    loss=loss,
)
trainer.train()
model.save_pretrained("./legal-embedding-model")

31.2.5 Practical Considerations

Theory and fine-tuning aside, three production gotchas catch nearly every team shipping their first embedding pipeline. The first is instruction prefixes, where forgetting "query: " or "passage: " can silently halve recall. The second is the asymmetry between query-time and document-time processing, and the third is batched re-embedding when models are versioned. We address them in that order.

Query and Document Prefixes

Many modern embedding models (E5, BGE, GTE) use instruction prefixes to distinguish between different input types. For example, the E5 model family expects queries to be prefixed with "query: " and documents with "passage: ". Forgetting these prefixes can significantly degrade retrieval performance. Always check the model documentation for required formatting conventions.

Sequence Length and Truncation

Every embedding model has a maximum sequence length, beyond which input is truncated. Older models like all-MiniLM-L6-v2 support only 256 tokens, while newer models such as nomic-embed-text-v1.5 and GTE-Qwen2 handle 8,192 or even 32,768 tokens. When your documents exceed the model's limit, you must chunk them before embedding, which is covered in detail in Section 31.6.

Batch Size and Throughput

Embedding throughput scales linearly with batch size up to the GPU memory limit. For production workloads, encode documents in large batches (256 to 1024) to maximize GPU utilization. For real-time query encoding, latency is more important than throughput, so smaller batch sizes (1 to 16) are typical. Consider using ONNX Runtime or TensorRT for optimized inference if you are serving embeddings at scale.

from sentence_transformers import SentenceTransformer
import numpy as np
model_small = SentenceTransformer("all-MiniLM-L6-v2") # 80MB, 384 dims
model_large = SentenceTransformer("all-mpnet-base-v2") # 420MB, 768 dims
print(f"Small: {model_small.get_sentence_embedding_dimension()} dims")
print(f"Large: {model_large.get_sentence_embedding_dimension()} dims")
# Quick test
test = "Machine learning is a subset of artificial intelligence."
print(f"Small shape: {model_small.encode(test).shape}")
print(f"Large shape: {model_large.encode(test).shape}")

documents = [
    "Python is a high-level programming language known for readability.",
    "Neural networks are computing systems inspired by biological brains.",
    "The Great Wall of China is over 13,000 miles long.",
    "Photosynthesis converts sunlight into chemical energy in plants.",
    "JavaScript is the most popular language for web development.",
    "Deep learning uses multiple layers of neural networks.",
    "The Amazon rainforest produces 20% of the world's oxygen.",
    "SQL is used for managing and querying relational databases.",
    "Climate change is causing rising sea levels worldwide.",
    "Transfer learning lets models trained on one task help with another.",
    "The human genome contains approximately 3 billion base pairs.",
    "Docker containers package applications with their dependencies.",
    "Coral reefs support 25% of all marine species.",
    "Gradient descent is the core optimization algorithm in deep learning.",
    "The International Space Station orbits Earth every 90 minutes.",
]
# TODO: Encode all documents with both models
embeddings_small = model_small.encode(documents)
embeddings_large = model_large.encode(documents)
print(f"Small embeddings: {embeddings_small.shape}")
print(f"Large embeddings: {embeddings_large.shape}")

import numpy as np
def semantic_search(query, model, doc_embeddings, documents, top_k=3):
    """Find the top-k most similar documents to the query."""
    # TODO: Encode the query, compute cosine similarity, return top-k
    query_emb = model.encode(query)
    norms = np.linalg.norm(doc_embeddings, axis=1) * np.linalg.norm(query_emb)
    scores = np.dot(doc_embeddings, query_emb) / norms
    top_idx = np.argsort(scores)[::-1][:top_k]
    return [(documents[i], scores[i], i) for i in top_idx]
# Test
queries = [
    "How do plants make food from sunlight?",
    "What programming language is best for web apps?",
    "Tell me about artificial intelligence techniques",
    "What are some environmental issues?",
    ]
for query in queries:
    print(f"\nQuery: {query}")
    for doc, score, _ in semantic_search(query, model_small, embeddings_small, documents):
        print(f" [{score:.3f}] {doc[:70]}...")

import pandas as pd
ground_truth = {
    "How do plants make food from sunlight?": [3],
    "What programming language is best for web apps?": [4],
    "Tell me about AI techniques": [1, 5, 9, 13],
    "What are environmental issues?": [8, 6, 12],
    }
def precision_at_k(retrieved_idx, relevant_idx, k=3):
    return len(set(retrieved_idx[:k]) & set(relevant_idx)) / k
results = []
for query, relevant in ground_truth.items():
    for name, mdl, embs in [("MiniLM", model_small, embeddings_small),
        ("MPNet", model_large, embeddings_large)]:
        hits = semantic_search(query, mdl, embs, documents, top_k=3)
        retrieved = [idx for _, _, idx in hits]
        p3 = precision_at_k(retrieved, relevant, k=3)
        results.append({"query": query[:40], "model": name, "P@3": p3})
        df = pd.DataFrame(results)
        print(df.to_string(index=False))
        print(f"\nAvg P@3 MiniLM: {df[df['model']=='MiniLM']['P@3'].mean():.3f}")
        print(f"Avg P@3 MPNet: {df[df['model']=='MPNet']['P@3'].mean():.3f}")

from sentence_transformers import SentenceTransformer
import numpy as np, pandas as pd
model_small = SentenceTransformer("all-MiniLM-L6-v2")
model_large = SentenceTransformer("all-mpnet-base-v2")
documents = [
    "Python is a high-level programming language known for readability.",
    "Neural networks are computing systems inspired by biological brains.",
    "The Great Wall of China is over 13,000 miles long.",
    "Photosynthesis converts sunlight into chemical energy in plants.",
    "JavaScript is the most popular language for web development.",
    "Deep learning uses multiple layers of neural networks.",
    "The Amazon rainforest produces 20% of the world's oxygen.",
    "SQL is used for managing and querying relational databases.",
    "Climate change is causing rising sea levels worldwide.",
    "Transfer learning lets models trained on one task help with another.",
    "The human genome contains approximately 3 billion base pairs.",
    "Docker containers package applications with their dependencies.",
    "Coral reefs support 25% of all marine species.",
    "Gradient descent is the core optimization algorithm in deep learning.",
    "The International Space Station orbits Earth every 90 minutes.",
    ]
embs_s = model_small.encode(documents)
embs_l = model_large.encode(documents)
def search(query, model, embs, docs, k=3):
    qe = model.encode(query)
    scores = np.dot(embs, qe) / (np.linalg.norm(embs, axis=1) * np.linalg.norm(qe))
    idx = np.argsort(scores)[::-1][:k]
    return [(docs[i], scores[i], i) for i in idx]
for q in ["How do plants make food?", "Best language for web apps?",
    "AI techniques", "Environmental issues?"]:
    print(f"\n{q}")
    for d, s, _ in search(q, model_small, embs_s, documents):
        print(f" [{s:.3f}] {d[:70]}")
        gt = {"How do plants make food?": [3], "Best language for web apps?": [4],
            "AI techniques": [1,5,9,13], "Environmental issues?": [8,6,12]}
def p_at_k(ret, rel, k=3):
    return len(set(ret[:k]) & set(rel)) / k
rows = []
for q, rel in gt.items():
    for nm, m, e in [("MiniLM",model_small,embs_s),("MPNet",model_large,embs_l)]:
        ri = [i for _,_,i in search(q, m, e, documents)]
        rows.append({"query":q[:35],"model":nm,"P@3":p_at_k(ri, rel)})
        df = pd.DataFrame(rows)
        print("\n" + df.to_string(index=False))
        for m in ["MiniLM","MPNet"]:
            print(f"Avg {m}: {df[df['model']==m]['P@3'].mean():.3f}")

from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

def search(query, model, embs, docs, k=3):
    qe = model.encode([query])
    sims = cosine_similarity(qe, embs)[0]
    idx = np.argsort(sims)[::-1][:k]
    return [(docs[i], sims[i], i) for i in idx]

import faiss, numpy as np

embs = doc_embeddings.astype("float32")
faiss.normalize_L2(embs)
index = faiss.IndexFlatIP(embs.shape[1])
index.add(embs)

qe = model.encode([query]).astype("float32"); faiss.normalize_L2(qe)
scores, idx = index.search(qe, top_k)