A Two-Stage Workflow to Detect Data Leakage.

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. Stage 1: Linear Leakage Check (Correlation & Single-Feature R2) 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.") Enter fullscreen mode Exit fullscreen mode How it Works Calculate the Correlation: The code starts by running a Pearson correlation between every feature in x_train and the target y_train. Correlation is a statistical metric that measures how closely two variables move together in a straight line, ranging from -1 to 1. We take the absolute value (.abs()) because a strong negative correlation is just as predictive as a strong positive one. Convert to R2: We then square those correlation values. Mathematically, squaring the Pearson correlation gives you the single-feature R2 (Coefficient of Determination). This tells us exactly what percentage of the variance in the target variable is explained by that one feature alone. Filter and Flag: Finally, we place the results in a DataFrame, sort them from highest to lowest, and flag any feature that explains more than 10% of the target's variance on its own. Why it is Effective 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. Stage 2: Non-Linear Leakage Check 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.") Enter fullscreen mode Exit fullscreen mode How it Works 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. Why it is Effective 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. The Output: === 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 Enter fullscreen mode Exit fullscreen mode The Takeaway 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.