# A Two-Stage Workflow to Detect Data Leakage. > Source: > Published: 2026-07-17 20:47:40+00:00 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.