Detect Train-Serving Skew

MLOps

Medium

Problem

Theory

One of the most frustrating scenarios in machine learning is when a model performs excellently during training and validation, but degrades mysteriously once deployed. Often, the culprit is train-serving skew, a mismatch between the data distribution the model was trained on and the data it encounters in production.

This skew can arise from many sources:

• Feature engineering differences: The training pipeline computes features slightly differently than the serving pipeline
• Data freshness: Production data reflects newer patterns not present in historical training data
• Population shift: The users or items in production differ from those in the training set
• Temporal drift: Seasonal or trend-based changes alter the input distribution over time

Detecting skew early is critical. If left unchecked, model performance silently degrades, predictions become unreliable, and downstream business decisions suffer.

Measuring Distribution Shift with PSI

The Population Stability Index (PSI) is a standard metric for quantifying how much a distribution has changed. Originally developed in credit risk modeling, it has become a go-to tool for monitoring ML feature distributions.

PSI compares two probability distributions (one from training as the reference, one from production as the target) by measuring the divergence bin by bin.

The Mathematical Definition

Given a feature discretized into n bins, let P_i^{train} be the proportion of training samples in bin i, and P_i^{serving} be the proportion of production samples in bin i.

The PSI is computed as:

PSI = \sum_{i=1}^{n} (P_i^{serving} - P_i^{train}) \times \ln\left(\frac{P_i^{serving}}{P_i^{train}}\right)

Each term in the sum captures both the magnitude of the difference (P_i^{serving} - P_i^{train}) and the relative change (\ln(P_i^{serving} / P_i^{train})). This makes PSI sensitive to both absolute and proportional shifts.

Interpreting PSI Values

Industry practice typically uses these thresholds:

• PSI < 0.1: No significant shift
• PSI between 0.1 and 0.25: Moderate shift, investigate
• PSI > 0.25: Significant shift, action required

A PSI above threshold indicates that the feature distribution in production has drifted meaningfully from training. This could signal that:

• The model learned patterns may no longer apply
• Feature engineering code may have diverged between pipelines
• A data quality issue has emerged

Handling Edge Cases

Zero-probability bins: If a bin has zero probability in either distribution, the logarithm and division become undefined. The standard fix is to add a small smoothing constant \epsilon (typically 10^{-10} to 10^{-6}) to all bin proportions:

P_i \leftarrow P_i + \epsilon

This ensures numerical stability while minimally affecting the PSI value for bins with non-zero counts.

Step-by-Step Computation

Given training proportions [P_1^{train}, P_2^{train}, ..., P_n^{train}] and serving proportions [P_1^{serving}, P_2^{serving}, ..., P_n^{serving}]:

Step 1: Add smoothing constant to prevent division by zero

P_i \leftarrow P_i + \epsilon \quad \text{for all } i

Step 2: For each bin i, compute the contribution:

\text{term}_i = (P_i^{serving} - P_i^{train}) \times \ln\left(\frac{P_i^{serving}}{P_i^{train}}\right)

Step 3: Sum all terms:

PSI = \sum_{i=1}^{n} \text{term}_i

Step 4: Compare against threshold to flag skew

A Concrete Example

Suppose a feature has 5 bins with these proportions:

• Bin 1: Training = 0.30, Serving = 0.25
• Bin 2: Training = 0.25, Serving = 0.20
• Bin 3: Training = 0.20, Serving = 0.25
• Bin 4: Training = 0.15, Serving = 0.20
• Bin 5: Training = 0.10, Serving = 0.10

Computing PSI:

• Bin 1: (0.25 - 0.30) \times \ln(0.25/0.30) = -0.05 \times (-0.182) = 0.0091
• Bin 2: (0.20 - 0.25) \times \ln(0.20/0.25) = -0.05 \times (-0.223) = 0.0112
• Bin 3: (0.25 - 0.20) \times \ln(0.25/0.20) = 0.05 \times 0.223 = 0.0112
• Bin 4: (0.20 - 0.15) \times \ln(0.20/0.15) = 0.05 \times 0.288 = 0.0144
• Bin 5: (0.10 - 0.10) \times \ln(0.10/0.10) = 0 \times 0 = 0

Total PSI = 0.0091 + 0.0112 + 0.0112 + 0.0144 + 0 = 0.0459

With a threshold of 0.1, this feature would not be flagged as skewed.

Where Train-Serving Skew Detection Shows Up

ML Model Monitoring: Production ML systems continuously compare incoming feature distributions against training baselines to detect drift before it impacts predictions.

A/B Testing: Before rolling out a new model, PSI helps verify that test and control groups have similar feature distributions, ensuring valid comparisons.

Data Quality Pipelines: ETL systems use PSI to catch upstream data changes that could propagate errors downstream.

Regulatory Compliance: In finance and healthcare, demonstrating that production data remains representative of training data is often a regulatory requirement.

Examples

Hints

1. Convert each pair of bin lists with np.asarray(..., dtype=float) + eps.
2. Compute one feature with np.sum((serving - train) * np.log(serving / train)).

Requirements

• Compute PSI independently for every feature
• Add eps before division and logarithms
• Mark values greater than or equal to the threshold as skewed
• Return {feature: {"psi": float, "skewed": bool}}

Constraints

• Training and serving dictionaries have the same feature keys
• Corresponding feature distributions have equal bin counts
• All bin proportions are nonnegative
• Use NumPy only

Starter Code

import numpy as np

def detect_skew(train_dist: dict, serving_dist: dict, threshold: float = 0.2, eps: float = 1e-10) -> dict:
    """
    Returns a dictionary of feature PSI scores and skew flags.
    """
    # Write code here
    pass

Test Cases