How to detect distribution shifts with Kolmogorov-Smirnov in Machine Learning

João Barros 12 de August de 2026 5 min read

This tutorial shows how to detect distribution shifts (data drift) in Machine Learning using the Kolmogorov-Smirnov test in Python. Knowing when features change over time is useful to keep models robust and avoid performance degradation. I will explain why the test is used, provide a practical code example and show how to interpret and act on the results.

Prerequisites

Python 3.8+ and libraries: pandas, scipy, matplotlib (install with pip).
A dataset split into two samples: reference (training) and current (production).
Basic knowledge of pandas and interpreting statistics.
For production, ideally have samples with representative sizes: e.g. 1k-10k for reference and 200-5k for current, depending on the case.

Step 1: Understand the purpose of the Kolmogorov-Smirnov test

The Kolmogorov-Smirnov (KS) test compares two samples to check if they come from the same continuous distribution. The result returns the KS statistic (the largest difference between the empirical distribution functions, in the interval [0,1]) and a p-value that quantifies the evidence against the null hypothesis that both samples come from the same distribution. The KS is nonparametric and sensitive to differences in location and shape of the distribution. This means it detects changes in mean, variance and even more subtle changes in the tails.

Step 2: Prepare example data

Create two samples: reference (training) and current (production) — with small changes to simulate drift. Keep column names identical. In the example below we use 1000 points in the reference and 500 in production; we changed the scale of 'income' to simulate drift and added some outliers and missing values to approximate real scenarios.

import numpy as np
import pandas as pd

np.random.seed(42)
# reference: normal and exponential distribution
reference = pd.DataFrame({
    'age': np.random.normal(40, 10, 1000),
    'income': np.random.exponential(40000, 1000)
})
# introduce some missing and outliers
reference.loc[::200, 'income'] = np.nan
reference.loc[5, 'age'] = 120  # outlier

# current: simulates drift in 'income' (scale change) and small change in age
current = pd.DataFrame({
    'age': np.random.normal(41, 11, 500),
    'income': np.random.exponential(60000, 500)
})
current.loc[::150, 'income'] = np.nan

Step 3: Implement the KS test per feature

We use scipy.stats.ks_2samp to compute the KS statistic and the p-value. Simple interpretation: low p-value (e.g.: < 0.05) indicates a significant difference between distributions. Note: with many features we have multiple tests; applying a correction (e.g. Bonferroni) reduces false positives. Also consider sample sizes: with large n small deviations can be statistically significant but without practical relevance.

from scipy.stats import ks_2samp

def detect_drift_ks(ref, cur, alpha=0.05, bonferroni=False):
    results = []
    numeric_cols = ref.select_dtypes(include=[np.number]).columns
    m = len(numeric_cols)
    for col in numeric_cols:
        a = ref[col].dropna()
        b = cur[col].dropna()
        stat, pvalue = ks_2samp(a, b)
        if bonferroni:
            pvalue_adj = min(pvalue * m, 1.0)
        else:
            pvalue_adj = pvalue
        drift = pvalue_adj < alpha
        results.append({
            'feature': col,
            'ks_stat': float(stat),
            'p_value': float(pvalue),
            'p_value_adj': float(pvalue_adj),
            'drift': bool(drift)
        })
    return pd.DataFrame(results)

results = detect_drift_ks(reference, current, alpha=0.05, bonferroni=True)
print(results)

Step 4: Visualize differences (example)

Visualizations help confirm what the test indicates. Plot histograms or KDE to compare distribution shapes between reference and current. If the KS indicates drift in 'income', the histogram typically shows a shift in the tail or increased dispersion. Include boxplots to see outliers and differences in the median.

import matplotlib.pyplot as plt

for col in ['age', 'income']:
    plt.figure(figsize=(6,3))
    plt.hist(reference[col].dropna(), bins=40, alpha=0.5, label='reference', density=True)
    plt.hist(current[col].dropna(), bins=40, alpha=0.5, label='current', density=True)
    plt.title(f'Comparação de distribuições: {col}')
    plt.legend()
    plt.tight_layout()
    plt.show()

Step 5: Handle detected drift

If the drift is significant, possible actions: re-train the model with updated data, apply re-calibration, use normalizations/transformations that are stable (e.g. quantile transformer trained on the reference), or implement continuous monitoring. Prioritize features with greater importance in the model and quantify the impact by re-evaluating performance metrics (AUC, RMSE, etc.). In many teams, a practical workflow is: 1) alert if 1-3 critical features have drift; 2) run an A/B test or retrain on a rolling window; 3) validate performance on a holdout validation.

Verify the result

Confirm the pipeline works by inspecting the results DataFrame: columns with 'drift' = True identify features with p-value < alpha. Validate visually with the histograms; if both indicate change, there is real drift. Test with different seeds/samples and check the stability of the test: sometimes small samples produce unstable p-values. Also consider magnitude metrics of drift (ks_stat) to prioritize interventions: values near 0.1 are weak, above 0.2-0.3 are already relevant changes in many contexts.

Conclusion

The Kolmogorov-Smirnov test is a simple and effective tool to detect distribution shifts in numerical features and help decide if a model needs maintenance. For production, automate this check, set thresholds per feature and combine with tests for categorical variables (e.g.: Chi-square) and with monitoring of model performance. Practical tip: start by monitoring the top 5-10 features with highest importance in the model and record metrics such as ks_stat and p_value over time to identify trends before acting.