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.")

How it Works

1. 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.
2. 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.
3. 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.")

How it Works

1. 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.

2. 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.

3. 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

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.