Build Embedding Drift Detection Using Distribution Divergence Metrics

Introduction

Engineers often discover that retrieval quality quietly degrades weeks after deployment — new documents shift the embedding distribution while the pipeline keeps returning results without any warning signal. Catching this drift requires comparing statistical snapshots of the embedding space over time using metrics that are numerically stable and easy to threshold. By the end of this lesson, you'll be able to build a sampler that captures distribution snapshots from a pgvector store, implement a Jensen-Shannon divergence detector that compares those snapshots against a stored baseline, and wire up a drift-alert pipeline that fires when distributional shift exceeds a configurable threshold.

Key Terminology

• Jensen-Shannon Divergence (JSD) — A symmetric, bounded divergence metric in [0, 1] that quantifies how different two probability distributions are by routing each through their pointwise mixture M = (P + Q) / 2 before computing KL; used here to compare normalized histograms of pairwise cosine similarities between a baseline and a current DistributionSnapshot.
• DistributionSnapshot — A dataclass that bundles a random sample of raw embeddings with precomputed summary statistics — centroid, variance, pairwise_cosine_mean, and pairwise_cosine_std — captured at a specific timestamp; the unit of comparison the drift detector operates on.
• Centroid Shift — The cosine distance between the mean embedding vectors (centroid fields) of two snapshots; a geometric signal that detects directional drift in the collection's aggregate semantic meaning, independent of JSD's shape-based measurement.
• TABLESAMPLE BERNOULLI — A PostgreSQL clause that gives each row an independent inclusion probability before a LIMIT cap, enabling probabilistic sampling without a full table scan; used in _fetch_sample to keep snapshot collection affordable regardless of collection size.
• Pairwise Cosine Similarity — The dot product of unit-normalized embedding pairs, averaged over a capped random subset of at most 1 000 pairs; captures the internal similarity structure of the embedding cloud and is histogrammed as the input distribution P or Q for JSD computation.
• KL Divergence (Kullback-Leibler) — The asymmetric base measure from which JSD is derived; undefined when a bin is zero in one distribution but non-zero in the other — the zero-bin failure mode that JSD eliminates by routing both distributions through their mixture M before computing KL.

Concepts

Why Snapshot Comparison Beats Per-Query Metrics

A single query's relevance score tells you whether that query degraded — it says nothing about the embedding space as a whole. Retrieval pipelines index new documents continuously, and no per-query signal captures whether the underlying vector cloud has restructured. The right unit of observation is the distribution itself: a statistical fingerprint of the embedding collection at a point in time. By storing a baseline fingerprint and periodically comparing it against a freshly drawn one, you can detect aggregate drift even when no individual query is obviously broken and before any user complaint surfaces.

The DistributionSnapshot dataclass materializes this fingerprint. Rather than persisting every vector, it summarizes the sample into the metrics that matter for drift: the centroid (aggregate direction), variance (spread around the centroid), and pairwise cosine similarity statistics (internal similarity structure). The detector then compares these summaries — not raw millions of vectors — making the approach practical for continuous scheduling (see Code Walkthrough).

Jensen-Shannon Divergence: Bounded, Symmetric, Always Defined

Comparing two distributions requires a metric that stays well-behaved when histograms don't overlap perfectly — a common situation in high-dimensional embedding spaces where two independent samples rarely share the exact same non-zero bins. KL divergence, the standard base measure, is undefined when a bin has zero probability in one distribution and non-zero in the other: log(p / 0) diverges. JSD eliminates the problem by routing both P and Q through their pointwise mixture M = (P + Q) / 2 before computing KL in each direction: JSD = (KL(P‖M) + KL(Q‖M)) / 2. Every bin of M is positive whenever at least one of P or Q is positive, so division by zero cannot occur.

JSD is also bounded to [0, 1], which makes threshold configuration interpretable across different embedding models and dimensionalities. The lesson establishes three operating bands — below 0.05 is routine sampling noise, above 0.15 warrants investigation, above 0.30 indicates severe shift — applied to the normalized histogram of pairwise cosine similarities within each snapshot. Histogramming the pairwise similarities, rather than the raw embedding coordinates, reduces a 1 536-dimensional problem to a 1-D distribution that JSD can compare directly.

Centroid Shift as a Complementary Geometric Signal

JSD measures changes in the shape of a distribution: whether spread, clustering, or tail behavior has changed. A complementary question is whether the collection has moved directionally: has the aggregate semantic center of all indexed documents migrated? Centroid shift answers this geometrically — the cosine distance between the centroid fields of two snapshots measures how far the mean embedding vector has rotated. A centroid that has shifted by more than a few degrees signals that new documents are pulling the cloud's center of mass toward a different semantic region, even when JSD remains moderate.

Together, the two signals provide complementary coverage. JSD catches restructuring — a bimodal split where a new topic cluster emerges while the mean stays roughly the same. Centroid shift catches directional drift — a topic influx where the whole collection moves toward a new subject area without necessarily changing its internal spread. Neither signal alone is sufficient; both together are difficult to fool.

Keeping Sampling Affordable at Scale

Two design choices in EmbeddingDistributionSampler bound computational cost regardless of collection size. The TABLESAMPLE BERNOULLI(5) clause in _fetch_sample gives each row an independent 5 % inclusion probability before the LIMIT $1 cap, so PostgreSQL performs a probabilistic scan — not a full table scan followed by a sort. Second, _pairwise_cosine_stats caps all-pairs computation at min(1000, len(embeddings)) randomly drawn pairs, keeping O(n²) pair enumeration out of the hot path entirely. Both choices mean each snapshot can be captured on a periodic schedule without blocking the pipeline — the prerequisite for continuous drift monitoring (see Code Walkthrough).

Code Walkthrough

Now that you understand how Jensen-Shannon divergence and centroid shift serve as complementary distributional signals, let's trace through a working implementation of the snapshot sampler that feeds both detectors.

The EmbeddingDistributionSampler connects to a pgvector database and returns a DistributionSnapshot — a dataclass that bundles raw embeddings alongside summary statistics: centroid, variance, and pairwise cosine similarity metrics. The _fetch_sample method issues a TABLESAMPLE BERNOULLI(5) query, giving each row an independent 5% chance of selection before capping at sample_size, so the database never performs a full table scan. Pairwise cosine similarity is computed on a capped random subset rather than all O(n²) combinations, keeping each snapshot affordable even at scale.

Code snippetpython
1import numpy as np
2from dataclasses import dataclass
3from datetime import datetime
4import asyncpg
5
6@dataclass
7class DistributionSnapshot:
8    timestamp: datetime
9    sample_size: int
10    centroid: np.ndarray
11    variance: float
12    pairwise_cosine_mean: float
13    pairwise_cosine_std: float
14    embeddings: np.ndarray
15
16class EmbeddingDistributionSampler:
17    def __init__(
18        self,
19        db_pool: asyncpg.Pool,
20        sample_size: int = 500,
21        embedding_dim: int = 1536,
22    ):
23        self.db_pool = db_pool
24        self.sample_size = sample_size
25        self.embedding_dim = embedding_dim
26
27    async def sample_and_summarize(self) -> DistributionSnapshot:
28        embeddings = await self._fetch_sample()
29        centroid = np.mean(embeddings, axis=0)
30        distances = np.linalg.norm(embeddings - centroid, axis=1)
31        variance = float(np.var(distances))
32        cosine_stats = self._pairwise_cosine_stats(embeddings)
33        return DistributionSnapshot(
34            timestamp=datetime.utcnow(),
35            sample_size=len(embeddings),
36            centroid=centroid,
37            variance=variance,
38            pairwise_cosine_mean=cosine_stats[0],
39            pairwise_cosine_std=cosine_stats[1],
40            embeddings=embeddings,
41        )
42
43    async def _fetch_sample(self) -> np.ndarray:
44        query = """
45            SELECT embedding
46            FROM document_embeddings
47            TABLESAMPLE BERNOULLI(5)
48            LIMIT $1
49        """
50        async with self.db_pool.acquire() as conn:
51            rows = await conn.fetch(query, self.sample_size)
52        return np.array([r["embedding"] for r in rows])
53
54    def _pairwise_cosine_stats(
55        self, embeddings: np.ndarray
56    ) -> tuple[float, float]:
57        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
58        normalized = embeddings / (norms + 1e-10)
59        n_pairs = min(1000, len(embeddings))
60        idx = np.random.choice(
61            len(embeddings), (n_pairs, 2), replace=True
62        )
63        sims = np.sum(
64            normalized[idx[:, 0]] * normalized[idx[:, 1]], axis=1
65        )
66        return float(np.mean(sims)), float(np.std(sims))

With a DistributionSnapshot in hand, a drift detector compares a stored baseline snapshot against a freshly drawn one. It bins each snapshot's pairwise cosine similarity values into normalized histograms, then computes JSD = (KL(P‖M) + KL(Q‖M)) / 2, where M = (P + Q) / 2. Because JSD is bounded between 0 and 1 and is always defined — no division-by-zero risk from empty bins — it maps directly to threshold-based alerting: values below 0.05 are routine noise, values above 0.15 warrant investigation, and values above 0.30 indicate severe distributional shift. The cosine distance between snapshot centroids adds an independent geometric check: a centroid that has shifted by more than a few degrees signals that the collection's aggregate meaning has changed, not just its spread.

Confirm that sample_and_summarize returns a DistributionSnapshot with embeddings.shape == (N, 1536) and that the JSD computed between two snapshots drawn from the same unchanged collection stays consistently below 0.05.

Do's and Don'ts

Having walked through the material above, the following Do's and Don'ts distill it into practice.

Do's

1. ✓Do use TABLESAMPLE BERNOULLI(5) with a LIMIT cap — this gives each row an independent 5% selection probability before capping at sample_size, so the database never performs a full table scan and snapshot cost stays predictable even as the document_embeddings table grows.
2. ✓Do compute pairwise cosine similarity on a capped random subset (e.g., min(1000, len(embeddings)) pairs) — computing all O(n²) combinations at sample_size=500 means 125,000 dot products per snapshot; the random-pair cap keeps each call affordable without meaningfully distorting the mean and std statistics that feed the JSD comparison.
3. ✓Do use JSD alongside centroid cosine distance as complementary signals — JSD on the binned cosine-similarity histograms detects spread and shape changes but can miss a uniform directional shift; the centroid-distance check catches exactly that case, so both signals must clear their respective thresholds before declaring the collection stable.

Don'ts

1. ✗Don't use KL divergence directly in place of JSD — KL(P‖M) is undefined when any bin in M has probability zero, which happens whenever the baseline and current snapshot histograms have non-overlapping support; JSD's symmetric midpoint mixture M = (P + Q) / 2 guarantees every bin is positive, keeping the detector numerically stable across arbitrary distributional shifts.
2. ✗Don't normalize embedding vectors before storing them in DistributionSnapshot.embeddings — the centroid and variance are computed from the raw distance of each embedding to the centroid using np.linalg.norm(embeddings - centroid, axis=1); pre-normalizing to unit vectors collapses magnitude information and makes the variance metric meaningless as a drift signal.
3. ✗Don't set a single fixed JSD threshold for all alert severities — the three-band model (below 0.05 = routine noise, 0.05–0.15 = monitor, above 0.30 = severe shift) is load-bearing: treating any non-zero JSD as an alert causes constant false positives from natural sampling variance, while a single high threshold misses gradual distributional drift that accumulates across multiple deployment cycles.