A Two-Stage Workflow to Detect Data Leakage.

A developer created a two-stage workflow to detect data leakage in machine learning models. The first stage uses fast correlation and single-feature R² to identify linear leaks, while the second stage trains a shallow decision tree to catch non-linear leaks via feature importance. The workflow flags features that explain more than 10% of the variance or dominate tree splits.

Data leakage is the silent killer of machine learning models. You finish training, see an incredible 98% accuracy on your test set, and feel like a genius, until the model hits production and completely fails. The most overlooked part of machine learning is the active prevention of data leakage. While robust cross-validation is a great way to verify if your model generalizes, it merely tells you that a problem exists. It doesn’t tell you where the leak is coming from. To solve this, I’ve created a standardized, two-stage workflow designed to identify the exact columns causing data leakage directly and immediately. Before you spend hours tuning hyperparameters, run your data through these two lightning-fast checks. Our first line of defense looks for direct, linear relationships. If a single column can perfectly or almost perfectly predict your target variable on its own, you likely have a leak—like an accidental "future" timestamp or an ID column that encodes the target label. ======================================== 1. Fast Correlation & Single-Feature R² ======================================== print "=== 1. Linear Leakage Check Correlation & R² ===" correlations = x train.corrwith pd.Series np.ravel y train , index=x train.index , numeric only=True .abs r2 approximations = correlations 2 leakage df = pd.DataFrame { 'Correlation |r| ': correlations, 'Single-Feature R²': r2 approximations } .sort values 'Single-Feature R²', ascending=False Flag features explaining 10% of variance high leakage = leakage df leakage df 'Single-Feature R²' 0.10 if len high leakage 0: print f"🔴 Found {len high leakage } potential leakage sources R² 10% :" print high leakage.head 20 else: print "✅ No linearly leaky features found." .abs because a strong negative correlation is just as predictive as a strong positive one.This step is computationally cheap. By relying on pandas' corrwith method, the operations are heavily vectorized in C, meaning it can scan thousands of columns in a fraction of a second. It immediately highlights glaring linear leaks without requiring you to train a single model. Linear correlation is great, but it misses complex, non-linear relationships. For example, a categorical ID mapped to integers might have a low linear correlation but still perfectly separate the target classes. To catch these, we move to Stage 2. ======================================== 2. Fast Non-Linear Leakage Check ======================================== print "\n=== 2. Non-Linear Leakage Check Tree Impurity ===" from sklearn.tree import DecisionTreeRegressor import pandas as pd dt = DecisionTreeRegressor max depth=5, random state=42 dt.fit x train, y train Built-in impurity importance costs 0 extra computation time tree importance = pd.DataFrame { 'Feature': x train.columns, 'Importance': dt.feature importances } .sort values 'Importance', ascending=False Features that alone dictate 10% of tree decisions top tree features = tree importance tree importance 'Importance' 0.10 if len top tree features 0: print f"🔴 Found {len top tree features } suspicious features dominating tree splits:" print top tree features.to string index=False else: print "✅ No single feature dominates the tree decisions." Train a Shallow Tree: We initialize a single DecisionTreeRegressor with a shallow depth max depth=5 . A decision tree works by continuously splitting the data into groups based on the feature that best separates the target variable. Extract Feature Importance: Once the tree is fitted, we pull the feature importances attribute. Scikit-learn calculates this based on impurity reduction—essentially tracking how much a specific feature helped the tree make clean splits. Filter and Flag: We map these importances back to the column names, sort them, and flag any feature that dictates more than 10% of the tree's overall decision-making process. A single decision tree restricted to a depth of 5 is incredibly fast to train. If a specific feature is a massive source of non-linear data leakage, the tree is going to be instantly drawn to it, splitting on that feature at the very top root nodes. Furthermore, extracting the built-in impurity importance requires zero extra computation time once the model is fitted. === 1. Linear Leakage Check Correlation & R² === 🔴 Found 6 potential leakage sources R² 10% : Correlation |r| \ target encoder apr drg code 0.586974 remainder total charges 0.578617 target encoder ccs diagnosis code 0.488559 target encoder ccs procedure code 0.480237 target encoder apr mdc code 0.376924 one hot encoder apr severity of illness code 4 0.347295 Single-Feature R² target encoder apr drg code 0.344538 remainder total charges 0.334797 target encoder ccs diagnosis code 0.238690 target encoder ccs procedure code 0.230628 target encoder apr mdc code 0.142072 one hot encoder apr severity of illness code 4 0.120614 === 2. Non-Linear Leakage Check Tree Impurity === 🔴 Found 2 suspicious features dominating tree splits: Feature Importance remainder total charges 0.752467 target encoder apr drg code 0.109771 By combining these two steps, you create an end-to-end, standardized workflow that catches both linear and non-linear data leakage before you commit to heavy model training. Embed this snippet at the top of your ML pipelines, and you'll never be caught off-guard by a leaky dataset in production again.