Detecting Distribution Shift in Production Machine Learning Systems
Machine learning systems deployed in real-world environments often experience performance degradation over time. One of the most common causes is distribution shift — where the statistical properties of incoming data diverge from the data used during training.
This is not a theoretical concern. In production, distribution shift is the default state. User behavior changes seasonally, upstream data pipelines break silently, sensor hardware gets swapped out, and third-party APIs alter their response schemas. If you are not actively monitoring for shift, you are flying blind.
This post covers the practical machinery of drift detection: statistical divergence measures with working code, embedding-space monitoring, a concrete medical imaging example, and the operational trade-offs around alerting thresholds.
Types of Distribution Shift
Before building monitors, it is worth being precise about what can shift.
| Covariate shift occurs when the input distribution P(X) changes while the conditional P(Y | X) remains stable. A classic example: a medical imaging model trained on images from Scanner A starts receiving images from Scanner B, which produces different brightness and contrast characteristics. The diagnostic relationship between image features and pathology has not changed, but the pixel-level statistics have. |
| Concept drift describes the case where P(Y | X) itself changes — the relationship between inputs and outputs evolves. In fraud detection, attackers continuously adapt their strategies, meaning the same transaction features map to different fraud probabilities over time. |
Label distribution shift (also called prior probability shift) refers to changes in the marginal P(Y). A disease screening model trained when prevalence was 2% may see prevalence rise to 8% during an outbreak, changing the base rate the model operates against.
In practice, these types co-occur. A recommender system experiencing seasonal changes may simultaneously see covariate shift (different item categories being browsed), concept drift (changing purchase intent), and label shift (different click-through rates). Your monitoring system should not assume the type of shift — it should detect the symptom and let the engineer diagnose the cause.
Statistical Divergence Measures
The foundation of any drift detection system is a way to quantify how different two distributions are. Two workhorses are KL divergence and the Population Stability Index (PSI).
KL Divergence
| KL divergence measures the information lost when approximating one distribution with another. It is asymmetric: D_KL(P | Q) ≠ D_KL(Q | P). For monitoring, we typically compute it from histograms of the feature values. |
Population Stability Index
PSI was originally developed in credit risk modeling and has become a standard metric for production ML monitoring. It is symmetric and has well-established interpretation thresholds:
- PSI < 0.1: no significant shift
- 0.1 ≤ PSI < 0.25: moderate shift — investigate
- PSI ≥ 0.25: significant shift — likely action required
import numpy as np
from typing import Optional
def compute_kl_divergence(
p: np.ndarray,
q: np.ndarray,
n_bins: int = 50,
epsilon: float = 1e-10,
bin_range: Optional[tuple[float, float]] = None,
) -> float:
"""
Compute KL divergence D_KL(P || Q) from two samples using histogram estimation.
Args:
p: samples from the reference distribution.
q: samples from the production distribution.
n_bins: number of histogram bins.
epsilon: smoothing constant to avoid log(0).
bin_range: optional (min, max) range for binning. If None, uses
the combined range of p and q.
Returns:
KL divergence (non-negative float; 0 = identical distributions).
"""
if bin_range is None:
bin_range = (min(p.min(), q.min()), max(p.max(), q.max()))
p_hist, bin_edges = np.histogram(p, bins=n_bins, range=bin_range, density=True)
q_hist, _ = np.histogram(q, bins=n_bins, range=bin_range, density=True)
# Smooth to prevent division by zero and log(0)
p_hist = p_hist + epsilon
q_hist = q_hist + epsilon
# Normalize to proper probability distributions
p_hist = p_hist / p_hist.sum()
q_hist = q_hist / q_hist.sum()
return float(np.sum(p_hist * np.log(p_hist / q_hist)))
def compute_psi(
reference: np.ndarray,
production: np.ndarray,
n_bins: int = 10,
epsilon: float = 1e-4,
) -> float:
"""
Compute Population Stability Index between reference and production distributions.
PSI = Σ (p_i - q_i) * ln(p_i / q_i)
where p_i and q_i are the proportions in bin i for reference and production
distributions respectively.
Args:
reference: samples from the training / reference distribution.
production: samples from the current production distribution.
n_bins: number of equal-width bins.
epsilon: smoothing constant for empty bins.
Returns:
PSI score (non-negative float).
"""
# Use reference distribution to define bin edges
bin_edges = np.percentile(reference, np.linspace(0, 100, n_bins + 1))
bin_edges[0] = -np.inf
bin_edges[-1] = np.inf
ref_counts = np.histogram(reference, bins=bin_edges)[0].astype(float)
prod_counts = np.histogram(production, bins=bin_edges)[0].astype(float)
# Convert to proportions
ref_proportions = ref_counts / ref_counts.sum() + epsilon
prod_proportions = prod_counts / prod_counts.sum() + epsilon
psi = np.sum(
(prod_proportions - ref_proportions) * np.log(prod_proportions / ref_proportions)
)
return float(psi)
A critical implementation detail: use percentile-based binning (quantile bins) derived from the reference distribution, not equal-width bins. Equal-width bins can produce empty or near-empty bins in the tails, making the metric unstable. Percentile binning ensures each bin has a meaningful number of reference samples, which gives you more reliable shift estimates in exactly the regions where shift matters most.
Embedding Drift Monitoring
For deep learning models, monitoring raw feature distributions is necessary but not sufficient. A model’s internal representations — its embeddings — can reveal distribution shift that is invisible at the feature level.
The idea is simple: compute embeddings for a reference dataset (e.g., a held-out validation set from training time), then compare production embeddings against this reference using cosine similarity.
┌─────────────────────────┐ ┌─────────────────────────┐
│ Training Data │ │ Production Data │
│ Distribution │ │ (live traffic) │
└───────────┬─────────────┘ └───────────┬─────────────┘
│ │
▼ ▼
┌─────────────────────────┐ ┌─────────────────────────┐
│ Compute Reference │ │ Compute Production │
│ Embeddings + Stats │ │ Embeddings │
│ (centroid, cov, etc.) │ │ │
└───────────┬─────────────┘ └───────────┬─────────────┘
│ │
▼ ▼
┌─────────┐ ┌─────────┐
│ Ref Stats│─────────► Monitor ◄───────────│Prod Stats│
└─────────┘ │ │ └─────────┘
│ │
┌───────┘ └───────┐
▼ ▼
┌────────────┐ ┌──────────────┐
│ Dashboard │ │ Alert System │
│ (Grafana) │ │ (PagerDuty) │
└────────────┘ └──────────────┘
import numpy as np
from dataclasses import dataclass
@dataclass
class EmbeddingReferenceStats:
"""Statistics computed from reference embeddings for drift comparison."""
centroid: np.ndarray # (D,) mean embedding
per_dim_mean: np.ndarray # (D,) per-dimension mean
per_dim_std: np.ndarray # (D,) per-dimension std
sample_embeddings: np.ndarray # (N, D) stored for distribution-level comparison
class EmbeddingDriftMonitor:
"""
Monitors embedding-level distribution shift by comparing production
embeddings against reference statistics.
"""
def __init__(self, reference_embeddings: np.ndarray):
"""
Args:
reference_embeddings: (N, D) array of embeddings from the
reference distribution (e.g., validation set).
"""
self.ref_stats = EmbeddingReferenceStats(
centroid=reference_embeddings.mean(axis=0),
per_dim_mean=reference_embeddings.mean(axis=0),
per_dim_std=reference_embeddings.std(axis=0),
sample_embeddings=reference_embeddings,
)
def cosine_similarity_to_centroid(self, production_embeddings: np.ndarray) -> float:
"""
Average cosine similarity between production embeddings and the
reference centroid. Values close to 1.0 = no drift; decreasing
values indicate the production distribution is moving away from
the reference center.
"""
centroid = self.ref_stats.centroid
centroid_norm = centroid / (np.linalg.norm(centroid) + 1e-10)
norms = np.linalg.norm(production_embeddings, axis=1, keepdims=True) + 1e-10
normalized = production_embeddings / norms
similarities = normalized @ centroid_norm
return float(np.mean(similarities))
def per_dimension_drift(
self, production_embeddings: np.ndarray, z_threshold: float = 3.0,
) -> dict:
"""
Check each embedding dimension for drift using z-score against
the reference mean and std.
Returns:
Dict with per-dimension z-scores and list of drifted dimensions.
"""
prod_means = production_embeddings.mean(axis=0)
# Standard error of the mean
n = len(production_embeddings)
z_scores = (
(prod_means - self.ref_stats.per_dim_mean)
/ (self.ref_stats.per_dim_std / np.sqrt(n) + 1e-10)
)
drifted_dims = np.where(np.abs(z_scores) > z_threshold)[0].tolist()
return {
"z_scores": z_scores,
"drifted_dimensions": drifted_dims,
"n_drifted": len(drifted_dims),
"max_abs_z": float(np.max(np.abs(z_scores))),
}
def compute_drift_report(self, production_embeddings: np.ndarray) -> dict:
"""Generate a complete drift report for a batch of production embeddings."""
cosine_sim = self.cosine_similarity_to_centroid(production_embeddings)
dim_drift = self.per_dimension_drift(production_embeddings)
# Also compute PSI on the embedding norm distribution —
# a surprisingly effective scalar summary of embedding drift
ref_norms = np.linalg.norm(self.ref_stats.sample_embeddings, axis=1)
prod_norms = np.linalg.norm(production_embeddings, axis=1)
norm_psi = compute_psi(ref_norms, prod_norms)
return {
"cosine_similarity_to_centroid": cosine_sim,
"embedding_norm_psi": norm_psi,
"per_dimension_drift": dim_drift,
"alert": cosine_sim < 0.85 or norm_psi > 0.25 or dim_drift["n_drifted"] > 5,
}
The embedding norm PSI is an underappreciated metric. When a model’s embeddings start producing vectors with noticeably different magnitudes, it often signals that inputs are falling outside the training manifold. This single scalar catches a surprising range of failure modes — from corrupted input features to upstream schema changes.
Concrete Example: Medical Imaging Pipeline
Let’s put this together with a realistic scenario. You have a convolutional model deployed to detect nodules in chest CT scans. The model was trained on images from a specific scanner fleet and validated to hit 0.94 AUC on a held-out test set. Six weeks after deployment, the radiology department installs new scanners. Nobody tells the ML team.
Here is what a monitoring pipeline for this system looks like:
import numpy as np
from typing import Any
class MedicalImagingDriftMonitor:
"""
Monitors pixel intensity distributions and embedding drift for a
deployed medical imaging model.
"""
def __init__(
self,
reference_pixel_stats: dict[str, np.ndarray],
reference_embeddings: np.ndarray,
psi_threshold: float = 0.20,
cosine_threshold: float = 0.88,
):
"""
Args:
reference_pixel_stats: dict with keys like 'intensity_mean',
'intensity_std', 'histogram' from reference images.
reference_embeddings: (N, D) model embeddings from reference images.
psi_threshold: PSI above this triggers an alert.
cosine_threshold: cosine sim below this triggers an alert.
"""
self.ref_pixel_stats = reference_pixel_stats
self.embedding_monitor = EmbeddingDriftMonitor(reference_embeddings)
self.psi_threshold = psi_threshold
self.cosine_threshold = cosine_threshold
self.alert_history: list[dict] = []
def compute_pixel_stats(self, images: np.ndarray) -> dict[str, Any]:
"""
Compute pixel-level statistics from a batch of images.
Args:
images: (B, H, W) or (B, H, W, C) array of images,
pixel values in [0, 1] or [0, 255].
"""
# Flatten spatial dimensions, keep batch
flat = images.reshape(images.shape[0], -1)
return {
"intensity_mean": flat.mean(axis=1), # per-image mean
"intensity_std": flat.std(axis=1), # per-image std
"intensity_p5": np.percentile(flat, 5, axis=1),
"intensity_p95": np.percentile(flat, 95, axis=1),
"global_pixel_values": flat.ravel(), # all pixels for histogram
}
def check_pixel_drift(self, production_images: np.ndarray) -> dict[str, Any]:
"""Check for pixel distribution drift between reference and production images."""
prod_stats = self.compute_pixel_stats(production_images)
results = {}
# PSI on per-image mean intensity
results["mean_intensity_psi"] = compute_psi(
self.ref_pixel_stats["intensity_mean"],
prod_stats["intensity_mean"],
)
# PSI on per-image contrast (std)
results["contrast_psi"] = compute_psi(
self.ref_pixel_stats["intensity_std"],
prod_stats["intensity_std"],
)
# KL divergence on global pixel value distribution
results["pixel_kl_divergence"] = compute_kl_divergence(
self.ref_pixel_stats["global_pixel_values"],
prod_stats["global_pixel_values"],
n_bins=100,
)
# PSI on the dynamic range (p95 - p5)
ref_range = self.ref_pixel_stats["intensity_p95"] - self.ref_pixel_stats["intensity_p5"]
prod_range = prod_stats["intensity_p95"] - prod_stats["intensity_p5"]
results["dynamic_range_psi"] = compute_psi(ref_range, prod_range)
results["alert"] = any(
v > self.psi_threshold
for k, v in results.items()
if k.endswith("_psi")
)
return results
def run_full_check(
self,
production_images: np.ndarray,
production_embeddings: np.ndarray,
batch_id: str,
) -> dict[str, Any]:
"""
Run complete drift check: pixel-level + embedding-level.
This is called periodically (e.g., every hour or every N predictions)
on a buffer of recent production data.
"""
pixel_report = self.check_pixel_drift(production_images)
embedding_report = self.embedding_monitor.compute_drift_report(production_embeddings)
combined_report = {
"batch_id": batch_id,
"pixel_drift": pixel_report,
"embedding_drift": embedding_report,
"overall_alert": pixel_report["alert"] or embedding_report["alert"],
}
if combined_report["overall_alert"]:
self.alert_history.append(combined_report)
return combined_report
When the new scanners are installed, this monitor would detect:
mean_intensity_psijumps from ~0.03 to ~0.35 — the new scanners produce brighter images on average.contrast_psirises to ~0.28 — the new scanners have a different dynamic range.- Embedding cosine similarity drops from 0.95 to 0.82 — the model’s internal representations are shifting because the pixel-level inputs look different.
The alert fires, the team investigates, and they discover the scanner change. The fix might be histogram equalization as a preprocessing step, retraining on data from the new scanners, or both.
Alerting Thresholds and Practical Trade-offs
Setting drift thresholds is fundamentally a precision-recall trade-off — and the right balance depends on the cost structure of your application.
| Application | False Negative Cost | False Positive Cost | Recommended Sensitivity |
|---|---|---|---|
| Medical diagnosis | Very high (patient harm) | Low (review queue) | High (sensitive thresholds) |
| Ad click prediction | Moderate (revenue loss) | Moderate (engineering time) | Medium |
| Content recommendation | Low (engagement dip) | High (alert fatigue) | Low (conservative thresholds) |
Practical guidelines:
Start with loose thresholds and tighten. It is better to miss early drift signals and tighten later than to burn out your on-call rotation with false alarms in the first week.
Use tiered alerting. Not every drift signal deserves a page. A reasonable setup:
def classify_alert_severity(drift_report: dict) -> str:
"""
Classify drift severity into tiers for routing to the right channel.
Tier 1 (page): strong evidence of drift + model performance degradation.
Tier 2 (ticket): drift detected but no confirmed performance impact yet.
Tier 3 (log): minor statistical fluctuation, log for trend analysis.
"""
pixel = drift_report["pixel_drift"]
emb = drift_report["embedding_drift"]
max_psi = max(
pixel.get("mean_intensity_psi", 0),
pixel.get("contrast_psi", 0),
emb.get("embedding_norm_psi", 0),
)
cosine_sim = emb.get("cosine_similarity_to_centroid", 1.0)
if max_psi > 0.5 or cosine_sim < 0.75:
return "tier_1_page"
elif max_psi > 0.25 or cosine_sim < 0.85:
return "tier_2_ticket"
elif max_psi > 0.10 or cosine_sim < 0.92:
return "tier_3_log"
else:
return "no_action"
Monitor the monitor. Track your drift metrics over time as time series. Many drift signals are seasonal or cyclical (e.g., user behavior on weekends vs weekdays). After a few weeks, you will have a baseline for what “normal” fluctuation looks like, and you can set thresholds relative to that baseline rather than using fixed constants.
Correlate drift with model performance. Drift detection is a leading indicator — it tells you inputs are changing — but not every input change degrades model performance. The highest-value monitoring setup pairs drift metrics with a delayed ground truth pipeline: once labels arrive (which may be hours or days after prediction time), compute model performance metrics and correlate them with the drift signals. This lets you learn which drift metrics actually predict performance degradation for your specific model, and calibrate thresholds accordingly.
Mitigation Strategies
Once drift is detected, the response depends on severity and root cause:
- Scheduled retraining: the simplest approach. Retrain on a rolling window of recent data at a fixed cadence (e.g., weekly). Works well for gradual drift but will not catch sudden shifts between retraining cycles.
- Triggered retraining: automatically kick off a retraining pipeline when drift exceeds a threshold. Requires a mature ML platform with automated training, validation, and deployment.
- Online learning: update model parameters continuously on incoming data. Powerful but operationally complex — you need safeguards against feedback loops and adversarial inputs.
- Domain adaptation: apply techniques like batch normalization recalibration or feature alignment to adapt a fixed model to the shifted distribution without full retraining.
- Input preprocessing: if the shift is in raw inputs (e.g., pixel intensities from a new scanner), preprocessing transforms like histogram matching can normalize inputs to look like the training distribution.
Conclusion
Distribution shift is not a failure mode — it is the steady state of any production ML system operating in the real world. The question is not whether your data will drift, but whether you will know about it when it does.
The monitoring stack described here — statistical divergence on raw features, embedding-space drift detection, tiered alerting, and correlation with model performance — is not theoretical. These are the primitives that production ML teams use to keep models reliable. The hard part is not implementing the metrics; it is building the operational discipline to act on them.
Mattia Gaggi is an applied machine learning engineer working on production ML systems.
Last updated 2026-04-13